from transformers import TextIteratorStreamer from transformers import AutoModelForCausalLM, AutoTokenizer from threading import Thread import gradio as gr class ChatbotService: def __init__(self, model_name="RajuKandasamy/tamillama_tiny_30m"): 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("
", "\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=512, 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"""Words: prevent, car, broken Features: Dialogue""", f"""சொற்கள்: தீர்வு""" ] chatbot_service = ChatbotService() with gr.Blocks() as demo: chatbot = gr.Chatbot().style(height=400) with gr.Row(): msg = gr.Textbox(placeholder="Type your message here...", inputs="text",outputs="text", label="Story Prompt:") run = gr.Button("Run") examples_dropdown = gr.Dropdown(choices=example_questions, 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()