# ===============================
# AmorCoder AI - Space de Producción
# ===============================
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel
import torch

# -------------------------------
# 1️⃣ Cargar modelo base y LoRA
# -------------------------------
MODEL_NAME = "codellama/CodeLlama-7b-hf"
LORA_PATH = "./lora_codellama"  # carpeta con tus pesos LoRA

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

print("Aplicando pesos LoRA...")
model = PeftModel.from_pretrained(model, LORA_PATH)

# Crear pipeline de generación
codegen = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=0,
    max_new_tokens=512,
    temperature=0.2,
    top_p=0.95
)

# -------------------------------
# 2️⃣ Función de generación de código
# -------------------------------
def generar_codigo(instruccion):
    prompt = f"# Instrucción:\n{instruccion}\n\n# Código:\n"
    salida = codegen(prompt)[0]["generated_text"]
    # Extraer solo el código generado después de la instrucción
    if "# Código:" in salida:
        return salida.split("# Código:")[1].strip()
    return salida.strip()

# -------------------------------
# 3️⃣ Interfaz Gradio
# -------------------------------
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()