Update app.py

app.py

@@ -1,16 +1,16 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Configuração de dispositivo
+# Detecta o dispositivo
 DEVICE = 0 if torch.cuda.is_available() else -1
 TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Carrega modelo A - LLaMA 2
+# Modelo A: Falcon 7B Instruct (sem autenticação)
 model_a = AutoModelForCausalLM.from_pretrained(
-    "meta-llama/Llama-2-7b-chat-hf",
+    "tiiuae/falcon-7b-instruct",
     torch_dtype=TORCH_DTYPE
 )
-tokenizer_a = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
+tokenizer_a = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
 
 pipe_a = pipeline(
     "text-generation",
@@ -21,12 +21,12 @@ pipe_a = pipeline(
     pad_token_id=tokenizer_a.eos_token_id
 )
 
-# Carrega modelo B - Falcon 7B Instruct (sem autenticação)
+# Modelo B: OpenAssistant Pythia 12B (sem autenticação)
 model_b = AutoModelForCausalLM.from_pretrained(
-    "tiiuae/falcon-7b-instruct",
+    "OpenAssistant/oasst-sft-1-pythia-12b",
     torch_dtype=TORCH_DTYPE
 )
-tokenizer_b = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
+tokenizer_b = AutoTokenizer.from_pretrained("OpenAssistant/oasst-sft-1-pythia-12b")
 
 pipe_b = pipeline(
     "text-generation",
@@ -37,24 +37,20 @@ pipe_b = pipeline(
     pad_token_id=tokenizer_b.eos_token_id
 )
 
-# Funções auxiliares para formatar o prompt
-def format_llama_prompt(user_input):
-    return f"[INST] <<SYS>>\nVocê é um assistente útil.\n<</SYS>>\n\n{user_input.strip()} [/INST]"
-
-def format_falcon_prompt(user_input):
+# Interface de prompt
+def format_prompt(user_input):
     return f"Responda em português: {user_input.strip()}"
 
-# Interface simples para testar os modelos
 if __name__ == "__main__":
     while True:
         prompt = input("\nDigite uma pergunta (ou 'sair'): ").strip()
         if prompt.lower() == "sair":
             break
 
-        print("\n=== Resposta do LLaMA 2 ===")
-        llama_response = pipe_a(format_llama_prompt(prompt), max_new_tokens=200)[0]['generated_text']
-        print(llama_response)
-
-        print("\n=== Resposta do Falcon 7B ===")
-        falcon_response = pipe_b(format_falcon_prompt(prompt), max_new_tokens=200)[0]['generated_text']
+        print("\n=== Resposta do Falcon ===")
+        falcon_response = pipe_a(format_prompt(prompt), max_new_tokens=200)[0]['generated_text']
         print(falcon_response)
+
+        print("\n=== Resposta do OpenAssistant ===")
+        oa_response = pipe_b(format_prompt(prompt), max_new_tokens=200)[0]['generated_text']
+        print(oa_response)