koko / app.py
viswazgummadi
done
a9833c6
import gradio as gr
import os
import PyPDF2
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Document
import re
from transformers import LlamaTokenizer, LlamaForCausalLM
print("started...........................................................")
model_path = "/mnt/data1/backup/viswaz/Project_K/huggingface_cache/Mistral-7B-Instruct-v0.2"
tokenizer = LlamaTokenizer.from_pretrained(model_path)
model = LlamaForCausalLM.from_pretrained(
model_path,
device_map="auto",
cache_dir="/mnt/data1/backup/viswaz/Project_K/huggingface_cache/",
)
print("stage1 ------------------ completed")
def generate_response(prompt, model):
encoded_input = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
model_inputs = encoded_input.to('cuda')
generated_ids = model.generate(**model_inputs,
max_new_tokens=150,
do_sample=True,
pad_token_id=tokenizer.eos_token_id)
return generated_ids
print("stage2 ------------------ completed")
intstructions_string = f""" you are a textbot that helps in finding answers to questions in the research papers, blogs,pdf's or any text context.
,make your answers more meaningful and short"
please answer the following question
"""
prompt_template_w_context = lambda context, question: f'''[INST] {intstructions_string}
{context}
Please answer to the following question. Use the context above if it is helpful.
{question}
[/INST]'''
print("stage3 ------------------ completed")
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.llm = None
Settings.chunk_size = 256
Settings.chunk_overlap = 25
top_k = 3
print("stage4 ------------------ completed")
def respond(files, question):
questions = question.split("\n")
documents = []
answers =""
texts_between_inst_and_eos = []
for file in files:
if file.endswith('.pdf'):
pdf_file_obj = open(file, 'rb')
pdf_reader = PyPDF2.PdfReader(pdf_file_obj)
text = ''
for page_num in range(len(pdf_reader.pages)):
page_obj = pdf_reader.pages[page_num]
text += page_obj.extract_text()
documents.append(Document(text=text)) # Join the extracted text
else:
with open(file, 'r') as f:
documents.append(Document(text=f.read())) # Join the text from the file
print("stage5 ------------------ completed")
joined_documents = "" # Initialize a variable to store the joined text
for doc in documents:
if "Member-only story" in doc.text:
continue
if "The Data Entrepreneurs" in doc.text:
continue
if " min read" in doc.text:
continue
joined_documents += doc.text # Join the text from the Document objects
print("stage6 ------------------ completed")
index = VectorStoreIndex.from_documents([Document(text=joined_documents)])
retriever = VectorIndexRetriever(
index=index,
similarity_top_k=top_k,
)
query_engine = RetrieverQueryEngine(
retriever=retriever,
node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.5)],
)
for each_question in questions:
response_each = query_engine.query(each_question)
context = "Context:\n"
for j in range(top_k):
context = context + response_each.source_nodes[j].text + "\n\n"
prompt = prompt_template_w_context(context, each_question)
outputs = generate_response(prompt, model)
decoded_output = tokenizer.batch_decode(outputs)[0]
inst_eos_texts = re.findall(r'\[\/INST\](.*?)\<\/s\>', decoded_output, re.DOTALL)
texts_between_inst_and_eos.extend(inst_eos_texts)
answers += ' '.join(texts_between_inst_and_eos)
print("stage7 ------------------ completed")
print(answers)
return answers
print("stage8 ------------------ completed")
inputs = [
gr.File(type="filepath", label="Upload PDF/Text Files",file_count="multiple"),
gr.Text(label="Enter your question here")
]
print("stage9 ------------------ completed")
output = gr.Text()
print("stage10 ------------------ completed")
interface = gr.Interface(respond, inputs=inputs, outputs=output, title="Question Answering System")
if __name__ == "__main__":
interface.launch()