Update app.py

app.py

@@ -1,58 +1,75 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 
-# Detecta o dispositivo
+# Configurações de memória
 DEVICE = 0 if torch.cuda.is_available() else -1
 TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# 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")
+# 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
+    }
+}
 
-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
-)
+# 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"]
 
-# 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
-)
+# Libera memória explicitamente
+def unload_model(pipe):
+    del pipe
+    torch.cuda.empty_cache()
 
 # Interface de prompt
 def format_prompt(user_input):
-    return f"Responda em português: {user_input.strip()}"
+    return f"Responda de forma clara e concisa: {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
 
-        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)
+        # 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)