Update app.py

app.py

@@ -1,21 +1,29 @@
 import gradio as gr
-from transformers import pipeline
+from llama_cpp import Llama
 
-# Carga un modelo un poco más potente pero rápido
-model = pipeline("text-generation", model="facebook/opt-350m")
+# Carga el modelo Gemma-2B-it desde llama_cpp
+llm = Llama.from_pretrained(
+    repo_id="google/gemma-2b-it-GGUF",  # Repositorio del modelo
+    filename="gemma-2b-it.gguf"         # Archivo del modelo
+)
 
 # Define una función para procesar las solicitudes
 def process_prompt(prompt):
-    response = model(prompt, max_length=100, num_return_sequences=1, do_sample=True)
-    return response[0]["generated_text"]
+    # Genera la salida del modelo
+    output = llm(
+        prompt,
+        max_tokens=512,  # Limita el número de tokens generados
+        echo=True        # Incluye el prompt en la salida
+    )
+    return output['choices'][0]['text']  # Retorna solo el texto generado
 
 # Configura la interfaz de API con Gradio
 interface = gr.Interface(
     fn=process_prompt, 
     inputs="text", 
     outputs="text",
-    title="Hugging Face Space API - OPT-350M",
-    description="Modelo ligero y rápido basado en OPT-350M para probar vulnerabilidades con RedTeamer.",
+    title="Hugging Face Space API - Gemma-2B-it",
+    description="Modelo basado en Gemma-2B-it para probar vulnerabilidades con RedTeamer.",
 )
 
 # Ejecuta la app