mandar100/chatbot_bloom3B

1#!/usr/bin/env python2# coding: utf-83 4# In[ ]:5 6 7from transformers import AutoModelForCausalLM, AutoTokenizer8import torch9import gradio as gr10import re11 12def cleaning_history_tuple(history):13    s=sum(history,())14    s=list(s)15    s2=""16    for i in s:17        i=re.sub("\n", '', i)18        i=re.sub("<p>", '', i)19        i=re.sub("</p>", '', i)20        s2=s2+i+'\n'21    return s222 23def ai_output(string1,string2):24    a1=len(string1)25    a2=len(string2)26    string3=string2[a1:]27    sub1="A:"28    sub2="User"29    #sub3="\n"30    try:31        try:32            idx1=string3.index(sub1)33            response=string3[:idx1]34            return response35            36        except:37            idx1=string3.index(sub2)38            response=string3[:idx1]39            return response40    except:41        return string342 43model4 = AutoModelForCausalLM.from_pretrained("bigscience/bloom-3b")44tokenizer4 = AutoTokenizer.from_pretrained("bigscience/bloom-3b")45 46def predict(input,initial_prompt, temperature=0.7,top_p=1,top_k=5,max_tokens=64,no_repeat_ngram_size=1,num_beams=6,do_sample=True, history=[]):47	48    s = cleaning_history_tuple(history)49    50    s = s+ "\n"+ "User: "+ input + "\n" + "Assistant: "51    s2=initial_prompt+" " + s52    53    input_ids = tokenizer4.encode(str(s2), return_tensors="pt")54    response = model4.generate(input_ids, min_length = 10,55                         max_new_tokens=int(max_tokens),56                         top_k=int(top_k),57                         top_p=float(top_p),58                         temperature=float(temperature),59                         no_repeat_ngram_size=int(no_repeat_ngram_size),60                         num_beams = int(num_beams),61                         do_sample = bool(do_sample),62                         )63    64  65    response2 = tokenizer4.decode(response[0])66    print("Response after decoding tokenizer: ",response2)67    print("\n\n")68    response3=ai_output(s2,response2)69    70    input="User: "+input71    response3="Assistant: "+ response372    history.append((input, response3))73 74    return history, history75 76#gr.Interface(fn=predict,title="BLOOM-3b",77#             inputs=["text","text","text","text","text","text","text","text","text",'state'],78#            79#             outputs=["chatbot",'state']).launch()80             81 82gr.Interface(inputs=[gr.Textbox(label="input", lines=1, value=""),83                     gr.Textbox(label="initial_prompt", lines=1, value=prompt),84                     gr.Textbox(label="temperature", lines=1, value=0.7),85                     gr.Textbox(label="top_p", lines=1, value=1),86                     gr.Textbox(label="top_k", lines=1, value=5),87                     gr.Textbox(label="max_tokens", lines=1, value=64),88                     gr.Textbox(label="no_repeat_ngram_size", lines=1, value=1),89                     gr.Textbox(label="num_beams", lines=1, value=6),90                     gr.Textbox(label="do_sample", lines=1, value="True"), 'state'],91             fn=predict, title="OPT-6.7B", outputs=["chatbot",'state']92 93             #inputs=["text","text","text","text","text","text","text","text","text",'state'],94 95             ).launch()            96 97