ValdeciRodrigues/Python-code-assistant

1# logic/generator.py2from transformers import AutoTokenizer, AutoModelForCausalLM3import torch4 5# Modelo poderoso, mas exige cuidado com limite de tokens6model_id = "stabilityai/stable-code-3b"7 8# Carregamento otimizado9tokenizer = AutoTokenizer.from_pretrained(model_id)10model = AutoModelForCausalLM.from_pretrained(11    model_id,12    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,13    device_map="auto"14)15 16def generate_code(prompt):17    try:18        # Instrução formatada para guiar o modelo19        formatted_prompt = f"# Escreva um código Python que faça o seguinte:\n# {prompt}\n"20 21        # Tokenização com truncamento seguro para evitar overflow22        inputs = tokenizer(23            formatted_prompt,24            return_tensors="pt",25            truncation=True,26            max_length=512  # entrada limitada para evitar travamento27        ).to(model.device)28 29        # Geração com finalização forçada via EOS token30        outputs = model.generate(31            **inputs,32            max_new_tokens=256,  # reduzido para caber nos limites de GPU33            do_sample=True,34            temperature=0.3,35            top_k=50,36            top_p=0.95,37            eos_token_id=tokenizer.eos_token_id  # 🚨 Essencial para evitar loop eterno38        )39 40        # Decodifica e retorna o texto limpo41        result = tokenizer.decode(outputs[0], skip_special_tokens=True)42        return result.strip()43 44    except Exception as e:45        return f"Erro ao gerar código: {str(e)}"46