import gradio as gr
import spaces
import llama_cpp
import llama_cpp.llama_tokenizer

import gradio as gr

llama = llama_cpp.Llama.from_pretrained(
    repo_id="jordigonzm/gemma-2b-it",
    filename="gemma-2b-it_v1p1-Q5_K_M.gguf",
    tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained("jordigonzm/gemma-2b-it"),
    verbose=False
)

model = "gpt-3.5-turbo"

def predict(message, history):
    messages = []
        
    # Añadir el nuevo mensaje del usuario.
    messages.append({"role": "user", "content": message})
    
    # Crear la respuesta del modelo.
    response = llama.create_chat_completion_openai_v1(
        model=model,
        messages=messages,
        stream=True
    )
    
    # Recopilar y devolver la respuesta.
    text = ""
    for chunk in response:
        content = chunk.choices[0].delta.content
        if content:
            text += content
            yield text

chat_interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your message here"),
    ],
    outputs="text",
    title="Chat with AI Model",
    description="Model Gemma-2b-it_v1p1 Q5_K_M",
    theme="soft",
)

chat_interface.launch()