from transformers import pipeline
import gradio as gr

from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM

# Configuración y carga del modelo personalizado
config = PeftConfig.from_pretrained("Juliofc/chaterapi_model")
base_model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
model = PeftModel.from_pretrained(base_model, "Juliofc/chaterapi_model")

# Inicializar el pipeline de generación de texto con el modelo cargado
chatbot = pipeline("text-generation", model=model, tokenizer="Juliofc/chaterapi_model")

def chat_with_bot(prompt):
    # Generar una respuesta del chatbot
    chat_response = chatbot(prompt, max_length=1000, num_return_sequences=1)
    # Devolver solo el texto de la respuesta
    return chat_response[0]['generated_text']

# Crear la interfaz de Gradio
interface = gr.Interface(
    fn=chat_with_bot,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Escribe algo para empezar a chatear..."),
    outputs="text",
    title="Chatbot Personalizado",
    description="Este es un chatbot interactivo que utiliza un modelo personalizado. Prueba a hablar con él escribiendo algo abajo."
)

# Ejecutar la interfaz de Gradio
if __name__ == "__main__":
    interface.launch()