import gradio as gr
import hf_transfer
from transformers import AutoModelForCausalLM, AutoTokenizer,StoppingCriteriaList,TextIteratorStreamer
from threading import Thread
import os
HFTOKEN=os.getenv("hftoken")

model = AutoModelForCausalLM.from_pretrained(
    "kubernetes-bad/chargen-v2",
    token = HFTOKEN
)
tknz=AutoTokenizer.from_pretrained("kubernetes-bad/chargen-v2",token=HFTOKEN)



"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})
    
   
    response = ""
    model_inputs = tokenizer.build_chat_input(history=messages, role='user').input_ids.to(
        next(model.parameters()).device)

    streamer = TextIteratorStreamer(tokenizer, timeout=600, skip_prompt=True)
    eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
                    tokenizer.get_command("<|observation|>")]
    generate_kwargs = {
        "input_ids": model_inputs,
        "streamer": streamer,
        "max_new_tokens": max_tokens,
        "do_sample": True,
        "top_p": top_p,
        "temperature": temperature,
        "stopping_criteria": StoppingCriteriaList([stop]),
        "repetition_penalty": 1,
        "eos_token_id": eos_token_id,
    }
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    for new_token in streamer:
        if new_token and '<|user|>' in new_token:
            new_token = new_token.split('<|user|>')[0]
        if new_token:
            history[-1][1] += new_token
        yield history

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
js_func = """
function refresh() {
    const url = new URL(window.location);

    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""
app = gr.ChatInterface(
    
    respond,
    js=js_func,
    additional_inputs=[
        gr.Textbox(value="You are a bot who generates perfect roleplaying charecters.", label="System message"),
        gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)


if __name__ == "__main__":
    app.launch()