| # logic/generator.py | |
| 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, | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| device_map="auto" | |
| ) | |
| 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" | |
| # 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, # reduzido para caber nos limites de GPU | |
| do_sample=True, | |
| temperature=0.3, | |
| top_k=50, | |
| 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)}" | |