import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch

# List of available premium models
premium_models = [
    "Qwen/Qwen2-1.5B-Instruct",
]

# Dictionary to cache loaded pipelines
pipeline_cache = {}

# Initial system prompt
default_system_prompt = "You are a ChatBuddy and chat with the user in a Human way."

def load_pipeline(model_name):
    if model_name not in pipeline_cache:
        print(f"Loading model: {model_name}")
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
        pipeline_cache[model_name] = pipe
    return pipeline_cache[model_name]

def chatbot(user_input, history, model_choice):
    pipe = load_pipeline(model_choice)

    # Prepare the chat messages
    messages = [{"role": "system", "content": default_system_prompt}]
    for pair in history:
        messages.append({"role": "user", "content": pair[0]})
        messages.append({"role": "assistant", "content": pair[1]})
    messages.append({"role": "user", "content": user_input})

    # Flatten into a prompt string
    prompt = ""
    for msg in messages:
        if msg["role"] == "system":
            prompt += f"<|system|> {msg['content']}\n"
        elif msg["role"] == "user":
            prompt += f"<|user|> {msg['content']}\n"
        elif msg["role"] == "assistant":
            prompt += f"<|assistant|> {msg['content']}\n"
    
    # Generate a response
    response = pipe(prompt, max_new_tokens=200, do_sample=True, top_p=0.95, temperature=0.7)[0]['generated_text']
    
    # Extract only the last assistant response
    split_res = response.split("<|assistant|>")
    final_response = split_res[-1].strip() if len(split_res) > 1 else response

    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": final_response})
    return "", history

with gr.Blocks() as demo:
    gr.Markdown("# 🤖 ChatBuddy - Advanced Chatbot with Selectable LLMs")

    with gr.Row():
        model_choice = gr.Dropdown(label="Select Model", choices=premium_models, value=premium_models[0])
        
    with gr.Row():
        model_choice = gr.Textbox(label="System Prompt", value=default_system_prompt)
    
    chatbot_ui = gr.Chatbot(type="messages")
    user_input = gr.Textbox(show_label=False, placeholder="Type your message and press Enter")
    clear_btn = gr.Button("Clear")

    state = gr.State([])

    user_input.submit(chatbot, [user_input, state, model_choice], [user_input, chatbot_ui])
    clear_btn.click(lambda: ([], ""), None, [chatbot_ui, state])

demo.launch()