import gradio as gr
from huggingface_hub import InferenceClient
import random

# Initialize the model
model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
client = InferenceClient(model)

def chat_response(prompt, history, seed, temp, tokens, top_p, rep_p):
    generate_kwargs = {
        "temperature": temp,
        "max_new_tokens": tokens,
        "top_p": top_p,
        "repetition_penalty": rep_p,
        "do_sample": True,
        "seed": seed,
    }

    # Include the chat history in the prompt
    formatted_prompt = "\n".join([f"Q: {user_prompt}\nA: {bot_response}" for user_prompt, bot_response in history]) + f"\nQ: {prompt}\nA:"
    
    output = ""
    
    # Generating text in streaming mode
    for response in client.text_generation(formatted_prompt, **generate_kwargs, stream=True):
        # Assuming response is directly a string or contains a message
        output += response  # Using response directly since it's a string

        # Yield the updated output for real-time display
        yield [(prompt, output)]

    # Append the full response to history after completion
    history.append((prompt, output))
    yield history  # Yielding the updated history

def clear_chat():
    return [], []  # Returning an empty history

# Gradio interface
with gr.Blocks() as app:
    gr.HTML("<center><h1>Chatbot</h1><h3>Ask your questions!</h3></center>")
    
    chat_box = gr.Chatbot(height=500)
    inp = gr.Textbox(label="Your Question", lines=5)
    btn = gr.Button("Ask")
    clear_btn = gr.Button("Clear")
    
    rand_seed = gr.Checkbox(label="Random Seed", value=True)
    seed_slider = gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=random.randint(1, 1111111111111111))
    tokens_slider = gr.Slider(label="Max new tokens", value=3840, minimum=0, maximum=8000)
    temp_slider = gr.Slider(label="Temperature", value=0.9, minimum=0.01, maximum=1.0)
    top_p_slider = gr.Slider(label="Top-P", value=0.9, minimum=0.01, maximum=1.0)
    rep_p_slider = gr.Slider(label="Repetition Penalty", value=1.0, minimum=0.1, maximum=2.0)

    # Handle button click to get chat response
    btn.click(
        lambda prompt: chat_response(prompt, [], seed_slider.value, temp_slider.value, tokens_slider.value, top_p_slider.value, rep_p_slider.value),
        inp,
        chat_box,
    )

    clear_btn.click(clear_chat, None, [inp, chat_box])

app.launch(share=True, auth=("admin", "0112358"))