File size: 2,243 Bytes
2df226a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from langchain.vectorstores import Chroma
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import pipeline
import torch
from langchain.llms import HuggingFacePipeline
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.chains import RetrievalQA
import gradio as gr
def chat(chat_history, user_input):
bot_response = qa_chain({"query": user_input})
bot_response = bot_response['result']
response = ""
for letter in ''.join(bot_response):
response += letter + ""
yield chat_history + [(user_input, response)]
checkpoint = "MBZUAI/LaMini-Flan-T5-783M"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
base_model = AutoModelForSeq2SeqLM.from_pretrained(
checkpoint,
device_map="auto",
torch_dtype = torch.float32)
embeddings = SentenceTransformerEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1")
db = Chroma(persist_directory="ipc_vector_data", embedding_function=embeddings)
pipe = pipeline(
'text2text-generation',
model = base_model,
tokenizer = tokenizer,
max_length = 512,
do_sample = True,
temperature = 0.3,
top_p= 0.95
)
local_llm = HuggingFacePipeline(pipeline=pipe)
qa_chain = RetrievalQA.from_chain_type(llm=local_llm,
chain_type='stuff',
retriever=db.as_retriever(search_type="similarity", search_kwargs={"k":2}),
return_source_documents=True,
)
with gr.Blocks() as gradioUI:
gr.Image('lawbot511.png')
with gr.Row():
chatbot = gr.Chatbot()
with gr.Row():
input_query = gr.TextArea(label='Input',show_copy_button=True)
with gr.Row():
with gr.Column():
submit_btn = gr.Button("Submit", variant="primary")
with gr.Column():
clear_input_btn = gr.Button("Clear Input")
with gr.Column():
clear_chat_btn = gr.Button("Clear Chat")
submit_btn.click(chat, [chatbot, input_query], chatbot)
submit_btn.click(lambda: gr.update(value=""), None, input_query, queue=False)
clear_input_btn.click(lambda: None, None, input_query, queue=False)
clear_chat_btn.click(lambda: None, None, chatbot, queue=False)
gradioUI.queue().launch()
|