Update app.py

app.py

@@ -6,7 +6,7 @@ from peft import PeftModel
 base_model_id = "openlm-research/open_llama_3b"
 adapter_path = "jalonso24/lora-lateblight-v3"
 
-# ✅ Load model fully onto CPU (no device_map)
+# ✅ Load model fully onto CPU (safe for Spaces)
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_id,
     torch_dtype=torch.float32
@@ -16,9 +16,23 @@ model = PeftModel.from_pretrained(base_model, adapter_path)
 tokenizer = AutoTokenizer.from_pretrained(base_model_id, use_fast=False)
 tokenizer.pad_token = tokenizer.eos_token
 
-def predict_risk(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
+# ✅ Few-shot prompt + user input
+def predict_risk(user_prompt):
+    few_shot = (
+        "Escenario: La variedad es Yungay, sembrada en octubre. "
+        "Hubo 2 mm de lluvia, temperatura máxima de 15°C y humedad promedio de 70%.\n"
+        "Respuesta esperada: Riesgo: Bajo\n\n"
+        "Escenario: La variedad es INIA-302 Amarilis, sembrada en noviembre. "
+        "La precipitación fue 18.4 mm, temperatura máxima 17.2°C, la mínima 6.1°C y la humedad promedio 84.12%.\n"
+        "Respuesta esperada:"
+    )
+    
+    # Append user prompt in same format
+    full_prompt = few_shot + "\n" + user_prompt.strip()
+
+    inputs = tokenizer(full_prompt, return_tensors="pt")
     inputs = {k: v.to("cpu") for k, v in inputs.items()}
+
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -27,21 +41,22 @@ def predict_risk(prompt):
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
+
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    return response.split("Respuesta esperada:")[-1].strip()  # remove echo
 
+# UI
 examples = [[
-    "La variedad es INIA-302 Amarilis, sembrada en noviembre.\n"
-    "La precipitación fue 18.4 mm, la temperatura máxima 17.2°C, "
-    "la mínima 6.1°C y la humedad promedio 84.12%.\n\n"
-    "¿Cuál es el riesgo de tizón tardío?"
+    "Escenario: La variedad es INIA-302 Amarilis, sembrada en noviembre. "
+    "La precipitación fue 18.4 mm, la temperatura máxima 17.2°C, la mínima 6.1°C "
+    "y la humedad promedio 84.12%.\nRespuesta esperada:"
 ]]
 
 gr.Interface(
     fn=predict_risk,
-    inputs=gr.Textbox(lines=6, label="Escenario de entrada (español)"),
-    outputs=gr.Textbox(lines=4, label="Respuesta del modelo"),
+    inputs=gr.Textbox(lines=6, label="Escenario en español"),
+    outputs=gr.Textbox(lines=3, label="Predicción del modelo"),
     title="🌱 Predicción de Tizón Tardío (OpenLLaMA 3B + LoRA)",
-    examples=examples,
-    description="Predice el riesgo de tizón tardío a partir de clima, variedad y fecha de siembra."
+    description="Predice el riesgo de tizón tardío (bajo, moderado o alto) según el clima y la variedad de papa.",
+    examples=examples
 ).launch()