import gradio as gr
from llama_cpp import Llama

# 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):
    # 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 - Gemma-2B-it",
    description="Modelo basado en Gemma-2B-it para probar vulnerabilidades con RedTeamer.",
)

# Ejecuta la app
if __name__ == "__main__":
    interface.launch(share=True)