Update logic/generator.py

logic/generator.py

@@ -2,7 +2,10 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
+# Modelo poderoso, mas exige cuidado com limite de tokens
 model_id = "stabilityai/stable-code-3b"
+
+# Carregamento otimizado
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -12,17 +15,31 @@ model = AutoModelForCausalLM.from_pretrained(
 
 def generate_code(prompt):
     try:
+        # Instrução formatada para guiar o modelo
         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)
+
+        # Tokenização com truncamento seguro para evitar overflow
+        inputs = tokenizer(
+            formatted_prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512  # entrada limitada para evitar travamento
+        ).to(model.device)
+
+        # Geração com finalização forçada via EOS token
         outputs = model.generate(
             **inputs,
-            max_new_tokens=256,
+            max_new_tokens=256,  # reduzido para caber nos limites de GPU
             do_sample=True,
             temperature=0.3,
             top_k=50,
-            top_p=0.95
+            top_p=0.95,
+            eos_token_id=tokenizer.eos_token_id  # 🚨 Essencial para evitar loop eterno
         )
+
+        # Decodifica e retorna o texto limpo
         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)}"