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

base_model_id = "openlm-research/open_llama_3b"
adapter_path = "jalonso24/lora-lateblight-v3"

# Load to CPU (safe on Spaces)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    torch_dtype=torch.float32
).to("cpu")

model = PeftModel.from_pretrained(base_model, adapter_path)
tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_fast=False)
tokenizer.pad_token = tokenizer.eos_token

# New structured prompt logic
def predict_risk(prompt):
    structured_prompt = (
        f"{prompt.strip()}\n\n"
        "Clasifica el riesgo de tizón tardío (elige una sola opción):\n"
        "- Bajo\n- Moderado\n- Alto\n\nRespuesta:"
    )

    inputs = tokenizer(structured_prompt, return_tensors="pt")
    inputs = {k: v.to("cpu") for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=10,
            temperature=0.5,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    result = response.split("Respuesta:")[-1].strip().split("\n")[0]
    return result

# Interface
examples = [[
    "Variedad: INIA-302 Amarilis\n"
    "Precipitación: 18.4 mm\n"
    "Temperatura Máxima: 17.2°C\n"
    "Temperatura Mínima: 6.1°C\n"
    "Humedad Promedio: 84.12%\n"
    "Mes de siembra: noviembre"
]]

gr.Interface(
    fn=predict_risk,
    inputs=gr.Textbox(lines=8, label="Escenario estructurado"),
    outputs=gr.Textbox(label="Predicción del modelo"),
    title="🌱 Riesgo de Tizón Tardío (LLaMA + LoRA)",
    description="Predice si el riesgo es Bajo, Moderado o Alto usando datos estructurados.",
    examples=examples
).launch()