from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import gradio as gr

class ChatbotService:
    def __init__(self, model_name="RajuKandasamy/marabutamil"):
        self.model = AutoModelForCausalLM.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.streamer = None

    def call(self, prompt):
        self.streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, timeout=5)
        prompt = prompt.replace("<br>", "\n")
        print(prompt)
        inputs = self.tokenizer(prompt, return_tensors="pt")
        print(inputs)
        kwargs = dict(input_ids=inputs["input_ids"], streamer=self.streamer, max_new_tokens=256, do_sample=True, top_p=0.8, top_k=500, temperature=0.001, repetition_penalty=1.4)
        thread = Thread(target=self.model.generate, kwargs=kwargs)
        thread.start()
        return ""


import gradio as gr

example_questions = [
    f"""பாடல்:
        நின்றன நின்றன நில்லாகும்""",
    f"""செல்வம் நிலையாமை
அறத்துப்பால்
பாடல்:
துகள்தீர் பெருஞ்செல்வம் தோன்றியக்கால் தொட்டுப்""",
    f"""பாடல்:
      கொங்குதேர் வாழ்க்கை அஞ்சிறைத் தும்பி
காமம் செப்பாது கண்டது மொழிமோ""",
    f"""வேதம் உரைத்தானும் வேதிய னாகிலன்
வேதம் உரைத்தானும் வேதா விளங்கிட"""
]

default_example = example_questions[3]
chatbot_service = ChatbotService()


with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    with gr.Row():
        msg = gr.Textbox(placeholder="Type your message here...", label="Venba first Stanza:", value=default_example)
        run = gr.Button("Run")
    examples_dropdown = gr.Dropdown(choices=example_questions, value=default_example, label="Select an example prompt")
    examples_dropdown.change(fn=lambda x: x, inputs=examples_dropdown, outputs=msg)

    clear = gr.Button("Clear")

    def user(question, user_message, history):
        if history == None:
            history = []
        user_message = question
        return "", history + [[user_message, None]]

    def bot(history):
        #print("Question: ", history[-1][0])
        chatbot_service.call(history[-1][0])
        history[-1][1] = ""
        for character in chatbot_service.streamer:
            print(character)
            history[-1][1] += character
            yield history

    run.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(bot, chatbot, chatbot)
    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue()
demo.launch()