Update logic/generator.py

logic/generator.py

@@ -1,8 +1,7 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Modelo leve compatível com Hugging Face Spaces
-model_id = "bigcode/starcoderbase"
+model_id = "stabilityai/stable-code-3b"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -11,10 +10,7 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 def generate_code(prompt):
-    # Formatação simples do prompt
     formatted_prompt = f"# Escreva um código Python que faça o seguinte:\n# {prompt}\n"
-
-    # Tokenização e envio para o modelo
     inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
@@ -24,7 +20,5 @@ def generate_code(prompt):
         top_k=50,
         top_p=0.95
     )
-
-    # Decodificação e retorno do resultado
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return result.strip()