from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os


MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"


tokenizer = AutoTokenizer.from_pretrained(MODEL)

model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.float16,
    device_map="auto"
)


def criar_arquivo(pedido):

    prompt = f"""
Você é um agente programador.

Crie somente código Python.

Tarefa:
{pedido}

Formato:
NOME_ARQUIVO:
conteúdo do arquivo
"""

    inputs = tokenizer(
        prompt,
        return_tensors="pt"
    ).to(model.device)


    output = model.generate(
        **inputs,
        max_new_tokens=1000
    )


    resposta = tokenizer.decode(
        output[0],
        skip_special_tokens=True
    )


    nome="programa.py"

    with open(nome,"w") as f:
        f.write(resposta)

    return nome