import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel
import torch

MODEL_NAME = "codellama/CodeLlama-7b-hf"
LORA_PATH = "lora_codellama"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    torch_dtype=torch.float16
)

model = PeftModel.from_pretrained(model, LORA_PATH)

codegen = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=0,
    max_new_tokens=512,
    temperature=0.2,
    top_p=0.95
)

def generar_codigo(instruccion):
    prompt = f"# Instrucción:\n{instruccion}\n\n# Código:\n"
    salida = codegen(prompt)[0]["generated_text"]
    if "# Código:" in salida:
        return salida.split("# Código:")[1].strip()
    return salida.strip()

iface = gr.Interface(
    fn=generar_codigo,
    inputs=gr.Textbox(lines=4, placeholder="Escribe tu instrucción aquí..."),
    outputs=gr.Textbox(lines=20, placeholder="Código generado..."),
    title="💙 AmorCoder AI - Programación Inteligente",
    description="Genera código real y completo usando CodeLlama 7B adaptado a tu estilo con LoRA."
)

iface.launch()