Spaces:
Runtime error
Runtime error
File size: 5,703 Bytes
0d5ea02 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
from pypdf import PdfReader
import torch
import PyPDF2
from io import BytesIO
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
import gradio as gr
import time
from langchain.memory import ConversationBufferMemory
from langchain.llms.huggingface_pipeline import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
from langchain.document_loaders import PyPDFDirectoryLoader
CHUNK_SIZE = 1000
# Using HuggingFaceEmbeddings with the chosen embedding model
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-mpnet-base-v2",model_kwargs = {"device": "cuda"})
# transformer model configuration
quant_config = BitsAndBytesConfig(
bnb_4bit_compute_dtype=torch.bfloat16
)
def load_llm():
model_id = "Deci/DeciLM-6b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id,
trust_remote_code=True,
device_map = "auto",
quantization_config=quant_config)
pipe = pipeline("text-generation",
model=model,
tokenizer=tokenizer,
temperature=0,
num_beams=5,
no_repeat_ngram_size=4,
early_stopping=True,
max_new_tokens=50,
)
llm = HuggingFacePipeline(pipeline=pipe)
return llm
def add_text(history, text):
if not text:
raise gr.Error('Enter text')
history = history + [(text, '')]
return history
def upload_file(file):
# file_path = [file.name for file in files]
print(type(file))
return file
def process_file(files):
# loader = PyPDFLoader(file_path= file.name)
# document = loader.load()
pdf_text = ""
for file in files:
# pdf_stream = BytesIO(file.name.content)
pdf = PyPDF2.PdfReader(file.name)
for page in pdf.pages:
pdf_text += page.extract_text()
# split into smaller chunks
text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=200)
splits = text_splitter.create_documents([pdf_text])
# create a FAISS vector store db
# embedd the chunks and store in the db
vectorstore_db = FAISS.from_documents(splits, embeddings)
#create a custom prompt
custom_prompt_template = """Given the uploaded files, generate a pecise answer to the question asked by the user.
If you don't know the answer, just say that you don't know, don't try to make up an answer.
Context= {context}
History = {history}
Question= {question}
Helpful Answer:
"""
prompt = PromptTemplate(template=custom_prompt_template, input_variables=["question", "context", "history"])
# set QA chain with memory
qa_chain_with_memory = RetrievalQA.from_chain_type(llm=load_llm(),
chain_type='stuff',
return_source_documents=True,
retriever=vectorstore_db.as_retriever(),
chain_type_kwargs={"verbose": True,
"prompt": prompt,
"memory": ConversationBufferMemory(
input_key="question",
memory_key="history",
return_messages=True) })
# get answers
return qa_chain_with_memory
def generate_bot_response(history,query, btn):
if not btn:
raise gr.Error(message='Upload a PDF')
qa_chain_with_memory = process_file(btn)
bot_response = qa_chain_with_memory({"query": query})
# return bot_response["result"]
for char in bot_response['result']:
history[-1][-1] += char
time.sleep(0.05)
yield history,''
with gr.Blocks() as demo:
with gr.Row():
with gr.Row():
chatbot = gr.Chatbot(label="DeciLM-6b-instruct bot", value=[], elem_id='chatbot')
with gr.Row():
file_output = gr.File(label="Your PDFs")
with gr.Column():
btn = gr.UploadButton("π Upload a PDF(s)", file_types=[".pdf"], file_count="multiple")
with gr.Column():
with gr.Column():
txt = gr.Text(show_label=False, placeholder="Enter question")
with gr.Column():
submit_btn = gr.Button('Ask')
# Event handler for uploading a PDF
btn.upload(fn=upload_file, inputs=[btn], outputs=[file_output])
submit_btn.click(
fn= add_text,
inputs=[chatbot, txt],
outputs=[chatbot],
queue=False
).success(
fn=generate_bot_response,
inputs=[chatbot, txt, btn],
outputs=[chatbot, txt]
).success(
fn=upload_file,
inputs=[btn],
outputs=[file_output]
)
if __name__ == "__main__":
demo.launch() |