Update app.py

app.py

@@ -2,14 +2,16 @@ import streamlit as st
 import requests
 from PIL import Image
 from io import BytesIO
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+import google.generativeai as genai
 
-# Configuración de la API
-API_URL = "https://api-inference.huggingface.co/models/riabayonaor/modelo_prediccion_enfermedades_pepinos"
+# Obtener la clave API de las variables de entorno
+gemini_api_key = os.getenv("GEMINI_API_KEY")
+genai.configure(api_key=gemini_api_key)
+model = genai.GenerativeModel('gemini-pro')
 
-# Cargar el modelo y el tokenizador de BLOOM
-tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom")
-model = AutoModelForCausalLM.from_pretrained("bigscience/bloom")
+# Configuración de la API de predicción de enfermedades
+API_URL = "https://api-inference.huggingface.co/models/riabayonaor/modelo_prediccion_enfermedades_pepinos"
 
 def query(image_bytes):
     response = requests.post(API_URL, data=image_bytes)
@@ -18,11 +20,9 @@ def query(image_bytes):
         return None
     return response.json()
 
-def bloom_query(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs["input_ids"], max_length=500, num_return_sequences=1)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+def gemini_query(prompt):
+    response = model.generate_content(prompt)
+    return response.text
 
 def main():
     st.set_page_config(page_title="Predicción de Enfermedades en Pepinos")
@@ -51,11 +51,11 @@ def main():
                 class_label = top_prediction['label']
                 st.write(f"Predicción principal: {class_label} con confianza {top_prediction['score']:.2f}")
 
-                # Usar la etiqueta principal para el modelo BLOOM
+                # Usar la etiqueta principal para el modelo Gemini
                 prompt = f"Esta enfermedad es {class_label}. Explica qué es y sugiere posibles insecticidas o soluciones en español."
 
-                # Llamar al modelo BLOOM
-                explanation = bloom_query(prompt)
+                # Llamar al modelo Gemini
+                explanation = gemini_query(prompt)
                 
                 st.write(f"Esta enfermedad es {class_label}:")
                 st.write(explanation)