import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import gradio as gr
from huggingface_hub import InferenceClient
from torch.cuda import is_available

from unsloth import FastLanguageModel
from transformers import TextIteratorStreamer
from threading import Thread

"""
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
"""

# client = InferenceClient()
class MyModel:
    def __init__(self):
        self.client = None
        self.current_model = ""
        self.tokenizer = None

    def respond(
        self,
        message,
        history: list[tuple[str, str]],
        model,
        system_message,
        max_tokens,
        temperature,
        min_p,
    ):
        if model != self.current_model or self.current_model is None:
            client, tokenizer = FastLanguageModel.from_pretrained(
                model_name = model,
                max_seq_length = 2048,
                dtype = None,
                load_in_4bit = True,
            )
            FastLanguageModel.for_inference(client) # Enable native 2x faster inference
            self.client = client
            self.tokenizer = tokenizer
            self.current_model = model
        
        text_streamer = TextIteratorStreamer(self.tokenizer, skip_prompt = True)

        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})

        inputs = self.tokenizer.apply_chat_template(
            messages,
            tokenize = True,
            add_generation_prompt = True, # Must add for generation
            return_tensors = "pt",
        ).to("cuda" if is_available() else "cpu")
        
        generation_kwargs = dict(input_ids=inputs, streamer=text_streamer, max_new_tokens=max_tokens, use_cache=True, temperature=temperature, min_p=min_p)
        thread = Thread(target=self.client.generate, kwargs=generation_kwargs)
        thread.start()

        response = ""

        for new_text in text_streamer:
            response += new_text
            yield response.strip("<|eot_id|>")

        # for message in client.chat_completion(
        #     messages,
        #     max_tokens=max_tokens,
        #     stream=True,
        #     temperature=temperature,
        #     top_p=top_p,
        #     model=model,
        # ):
        #     token = message.choices[0].delta.content

        #     response += token
        #     yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
my_model = MyModel()
model_choices = [
    "lab2-as/lora_model_gguf",
    "lab2-as/lora_model",
]
demo = gr.ChatInterface(
    my_model.respond,
    additional_inputs=[
        gr.Dropdown(choices=model_choices, value=model_choices[0], label="Select Model"),
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=128, 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="Min-p (nucleus sampling)",
        ),
    ],
)


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