iamkhadke/GeneralChatBot

1import gradio as gr2import torch3from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer4import time5import numpy as np6from torch.nn import functional as F7import os8from threading import Thread9 10print(f"Starting to load the model to memory")11#m = AutoModelForCausalLM.from_pretrained(12#    "stabilityai/stablelm-tuned-alpha-3b", torch_dtype=torch.float16).cuda()13#tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-7b")14from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig15 16quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)17tok = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-3b", device_map="auto", load_in_8bit=True, torch_dtype=torch.float16 )18m = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-3b", device_map= "auto", quantization_config=quantization_config,19                                        offload_folder="./")20# generator = pipeline('text-generation', model=m, tokenizer=tok, device=1)21print(f"Sucessfully loaded the model to the memory")22 23start_message = """<|SYSTEM|># StableAssistant24- StableAssistant is A helpful and harmless Open Source AI Language Model developed by Stability and CarperAI.25- StableAssistant is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.26- StableAssistant is more than just an information source, StableAssistant is also able to write poetry, short stories, and make jokes.27- StableAssistant will refuse to participate in anything that could harm a human."""28 29 30class StopOnTokens(StoppingCriteria):31    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:32        stop_ids = [50278, 50279, 50277, 1, 0]33        for stop_id in stop_ids:34            if input_ids[0][-1] == stop_id:35                return True36        return False37 38 39def user(message, history):40    # Append the user's message to the conversation history41    return "", history + [[message, ""]]42 43 44def chat(curr_system_message, history):45    # Initialize a StopOnTokens object46    stop = StopOnTokens()47 48    # Construct the input message string for the model by concatenating the current system message and conversation history49    messages = curr_system_message + \50        "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]])51                for item in history])52 53    # Tokenize the messages string54    model_inputs = tok([messages], return_tensors="pt")55    streamer = TextIteratorStreamer(56        tok, skip_prompt=True, skip_special_tokens=True)57    generate_kwargs = dict(58        model_inputs,59        streamer=streamer,60        max_new_tokens=1024,61        do_sample=True,62        top_p=0.95,63        top_k=1000,64        temperature=1.0,65        num_beams=1,66        stopping_criteria=StoppingCriteriaList([stop])67    )68    t = Thread(target=m.generate, kwargs=generate_kwargs)69    t.start()70 71    # print(history)72    # Initialize an empty string to store the generated text73    partial_text = ""74    for new_text in streamer:75        # print(new_text)76        partial_text += new_text77        history[-1][1] = partial_text78        # Yield an empty string to cleanup the message textbox and the updated conversation history79        yield history80    return partial_text81 82 83with gr.Blocks() as demo:84    # history = gr.State([])85    gr.Markdown("## Start your Chat")86    gr.HTML('''<center>Enter your text. Inference might be slower. \n<a href="https://www.linkedin.com/in/khadke-chetan/">Follow me</a></center>''')87    chatbot = gr.Chatbot().style(height=500)88    with gr.Row():89        with gr.Column():90            msg = gr.Textbox(label="Chat Message Box", placeholder="Chat Message Box",91                             show_label=False).style(container=False)92        with gr.Column():93            with gr.Row():94                submit = gr.Button("Submit")95                stop = gr.Button("Stop")96                clear = gr.Button("Clear")97    system_msg = gr.Textbox(98        start_message, label="System Message", interactive=False, visible=False)99 100    submit_event = msg.submit(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False).then(101        fn=chat, inputs=[system_msg, chatbot], outputs=[chatbot], queue=True)102    submit_click_event = submit.click(fn=user, inputs=[msg, chatbot], outputs=[msg, chatbot], queue=False).then(103        fn=chat, inputs=[system_msg, chatbot], outputs=[chatbot], queue=True)104    stop.click(fn=None, inputs=None, outputs=None, cancels=[105               submit_event, submit_click_event], queue=False)106    clear.click(lambda: None, None, [chatbot], queue=False)107 108demo.queue(max_size=32, concurrency_count=2)109demo.launch()