Update logic/generator.py

logic/generator.py

@@ -1,3 +1,4 @@
+# logic/generator.py
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
@@ -10,15 +11,18 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 def generate_code(prompt):
-    formatted_prompt = f"# Escreva um código Python que faça o seguinte:\n# {prompt}\n"
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.3,
-        top_k=50,
-        top_p=0.95
-    )
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return result.strip()
+    try:
+        formatted_prompt = f"# Escreva um código Python que faça o seguinte:\n# {prompt}\n"
+        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=True,
+            temperature=0.3,
+            top_k=50,
+            top_p=0.95
+        )
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return result.strip()
+    except Exception as e:
+        return f"Erro ao gerar código: {str(e)}"