Spaces:
Sleeping
Sleeping
import gradio as gr | |
import pandas as pd | |
from langchain.embeddings import HuggingFaceEmbeddings | |
from langchain.vectorstores import faiss | |
from langchain_community.llms import HuggingFaceHub | |
from langchain.chains import ConversationalRetrievalChain, LLMChain | |
from langchain_community.document_loaders import PyPDFLoader | |
from langchain.text_splitter import RecursiveCharacterTextSplitter | |
from langchain import vectorstores | |
from langchain.prompts import PromptTemplate | |
from langchain.memory import ConversationBufferMemory | |
import os | |
import re | |
api_token = os.environ.get('HUGGINGFACEHUB_API_TOKEN') | |
memory = ConversationBufferMemory( | |
memory_key="chat_history", | |
return_messages=True | |
) | |
model = HuggingFaceHub( | |
huggingfacehub_api_token=api_token, | |
repo_id="mistralai/Mistral-7B-Instruct-v0.2", | |
task="conversational", | |
model_kwargs={"temperature": 0.8, "max_length": 1000}, | |
) | |
template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use three sentences maximum. Keep the answer as concise as possible. Always say "thanks for asking!" at the end of the answer. | |
{context} | |
Question: {question} | |
Helpful Answer:""" | |
QA_CHAIN_PROMPT = PromptTemplate.from_template(template) | |
def load_db(file, k): | |
# load documents | |
loader = PyPDFLoader(file) | |
documents = loader.load() | |
# split documents | |
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150) | |
docs = text_splitter.split_documents(documents) | |
# define embedding | |
embeddings = HuggingFaceEmbeddings() | |
# create vector database from data | |
db = vectorstores.FAISS.from_documents(docs, embeddings) | |
# define retriever | |
retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": k}) | |
# create a chatbot chain. Memory is managed externally. | |
question_generator_chain = LLMChain(llm=model, prompt=QA_CHAIN_PROMPT) | |
qa = ConversationalRetrievalChain.from_llm( | |
llm=model, | |
chain_type="stuff", | |
retriever=retriever, | |
return_source_documents=True, | |
return_generated_question=True, | |
memory=memory, | |
) | |
return qa | |
def chat(input_text, pdf_file): | |
qa = load_db(pdf_file, 3) | |
print("MEMORY") | |
print(memory) | |
if not memory.history: | |
# If no previous conversation, start with a greeting | |
response = qa.invoke({"question": "Hi, how can I help you today?", "chat_history": []}) | |
memory.update(response["chat_history"]) | |
response = qa.invoke({"question": input_text, "chat_history": memory.history}) | |
# Extracting the helpful answer from the response | |
match = re.search(r'Helpful Answer:(.*)', response['answer']) | |
if match: | |
helpful_answer = match.group(1).strip() | |
else: | |
helpful_answer = "No helpful answer found." | |
# Update the chat history | |
memory.update([(input_text, helpful_answer)]) | |
return helpful_answer | |
iface = gr.Interface(fn=chat, inputs=["text", "file"], outputs="text") | |
iface.launch(share=True) | |