|
import gradio as gr |
|
from langchain.document_loaders import PDFMinerLoader, PyMuPDFLoader |
|
from langchain.text_splitter import CharacterTextSplitter |
|
import chromadb |
|
import chromadb.config |
|
from chromadb.config import Settings |
|
from transformers import T5ForConditionalGeneration, AutoTokenizer |
|
import torch |
|
import uuid |
|
from sentence_transformers import SentenceTransformer |
|
import os |
|
|
|
|
|
model_name = 'google/flan-t5-base' |
|
model = T5ForConditionalGeneration.from_pretrained(model_name) |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
|
|
|
ST_name = 'sentence-transformers/sentence-t5-base' |
|
st_model = SentenceTransformer(ST_name) |
|
|
|
client = chromadb.Client() |
|
collection = client.create_collection("my_database") |
|
|
|
|
|
def get_context(query_text): |
|
|
|
query_emb = st_model.encode(query_text) |
|
query_response = collection.query(query_embeddings=query_emb.tolist(), n_results=4) |
|
context = query_response['documents'][0][0] |
|
context = context.replace('\n', ' ').replace(' ', ' ') |
|
return context |
|
|
|
def local_query(query, context): |
|
t5query = """Using the available context, please answer the question. |
|
If you are not sure please say I don't know. |
|
Context: {} |
|
Question: {} |
|
""".format(context, query) |
|
|
|
inputs = tokenizer(t5query, return_tensors="pt") |
|
|
|
outputs = model.generate(**inputs, max_new_tokens=20) |
|
|
|
return tokenizer.batch_decode(outputs, skip_special_tokens=True) |
|
|
|
|
|
|
|
|
|
|
|
def run_query(btn, history, query): |
|
|
|
context = get_context(query) |
|
result = local_query(query, context) |
|
history.append((query, str(result[0]))) |
|
return history, "" |
|
|
|
|
|
|
|
def upload_pdf(file): |
|
try: |
|
if file is not None: |
|
|
|
global collection |
|
|
|
file_name = file.name |
|
|
|
loader = PDFMinerLoader(file_name) |
|
doc = loader.load() |
|
|
|
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0) |
|
texts = text_splitter.split_documents(doc) |
|
|
|
texts = [i.page_content for i in texts] |
|
|
|
doc_emb = st_model.encode(texts) |
|
doc_emb = doc_emb.tolist() |
|
|
|
ids = [str(uuid.uuid1()) for _ in doc_emb] |
|
|
|
|
|
collection.add( |
|
embeddings=doc_emb, |
|
documents=texts, |
|
ids=ids |
|
) |
|
|
|
|
|
return 'Successfully uploaded!' |
|
else: |
|
return "No file uploaded." |
|
|
|
except Exception as e: |
|
return f"An error occurred: {e}" |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
|
btn = gr.UploadButton("Upload a PDF", file_types=[".pdf"]) |
|
output = gr.Textbox(label="Output Box", style={"height": "100px", "margin-top": "20px"}) |
|
chatbot = gr.Chatbot(height=240) |
|
|
|
with gr.Row(): |
|
with gr.Column(scale=0.70): |
|
txt = gr.Textbox( |
|
show_label=False, |
|
placeholder="Enter a question", |
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
gr.close_all() |
|
|
|
demo.queue().launch() |
|
|
|
|
