Update app.py

app.py

@@ -1,75 +1,58 @@
 import torch
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Configurações de memória
+# Detecta o dispositivo
 DEVICE = 0 if torch.cuda.is_available() else -1
 TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Modelos otimizados para 16GB
-MODEL_CONFIG = {
-    "Modelo A": {
-        "name": "pierreguillain/gpt2-small-portuguese",
-        "max_tokens": 150
-    },
-    "Modelo B": {
-        "name": "pierreguillain/gpt-neo-125m-portuguese",
-        "max_tokens": 150
-    }
-}
+# Modelo A: Falcon 7B Instruct (sem autenticação)
+model_a = AutoModelForCausalLM.from_pretrained(
+    "tiiuae/falcon-7b-instruct",
+    torch_dtype=TORCH_DTYPE
+)
+tokenizer_a = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
 
-# Carrega apenas um modelo por vez
-def load_model(model_name):
-    config = MODEL_CONFIG[model_name]
-    model = AutoModelForCausalLM.from_pretrained(
-        config["name"],
-        torch_dtype=TORCH_DTYPE,
-        low_cpu_mem_usage=True  # Reduz consumo de memória
-    )
-    tokenizer = AutoTokenizer.from_pretrained(config["name"])
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=DEVICE,
-        return_full_text=False,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    return pipe, config["max_tokens"]
+pipe_a = pipeline(
+    "text-generation",
+    model=model_a,
+    tokenizer=tokenizer_a,
+    device=DEVICE,
+    return_full_text=False,
+    pad_token_id=tokenizer_a.eos_token_id
+)
 
-# Libera memória explicitamente
-def unload_model(pipe):
-    del pipe
-    torch.cuda.empty_cache()
+# Modelo B: OpenAssistant Pythia 12B (sem autenticação)
+model_b = AutoModelForCausalLM.from_pretrained(
+    "OpenAssistant/oasst-sft-1-pythia-12b",
+    torch_dtype=TORCH_DTYPE
+)
+tokenizer_b = AutoTokenizer.from_pretrained("OpenAssistant/oasst-sft-1-pythia-12b")
+
+pipe_b = pipeline(
+    "text-generation",
+    model=model_b,
+    tokenizer=tokenizer_b,
+    device=DEVICE,
+    return_full_text=False,
+    pad_token_id=tokenizer_b.eos_token_id
+)
 
 # Interface de prompt
 def format_prompt(user_input):
-    return f"Responda de forma clara e concisa: {user_input.strip()}"
+    return f"Responda em português: {user_input.strip()}"
 
 if __name__ == "__main__":
-    print("Sistema otimizado para 16GB de RAM\n")
-    
     while True:
         prompt = input("\nDigite uma pergunta (ou 'sair'): ").strip()
         if prompt.lower() == "sair":
             break
 
-        # Processa um modelo por vez
-        for model_name in MODEL_CONFIG:
-            try:
-                print(f"\n=== Carregando {model_name} ===")
-                pipe, max_tokens = load_model(model_name)
-                
-                print(f"\n=== Resposta do {model_name} ===")
-                response = pipe(
-                    format_prompt(prompt),
-                    max_new_tokens=max_tokens,
-                    temperature=0.7,
-                    top_p=0.9
-                )[0]['generated_text'].strip()
-                
-                print(response)
-                unload_model(pipe)
-                
-            except Exception as e:
-                print(f"Erro no {model_name}: {str(e)}")
-                unload_model(pipe)
+        print("\n=== Resposta do Falcon ===")
+        with torch.no_grad():
+            falcon_response = pipe_a(format_prompt(prompt), max_new_tokens=100)[0]['generated_text']
+        print(falcon_response)
+
+        print("\n=== Resposta do OpenAssistant ===")
+        with torch.no_grad():
+            oa_response = pipe_b(format_prompt(prompt), max_new_tokens=100)[0]['generated_text']
+        print(oa_response)