Update app.py

app.py

@@ -1,11 +1,39 @@
 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
 
 
 import gradio as gr
 
+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 aren't 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(query):
+    context = get_context(query)
+    result = local_query(query, context)
+    return result
+
+
+    
 def upload_pdf(file):
     # Save the uploaded file
     file_name = file.name
@@ -20,7 +48,42 @@ def upload_pdf(file):
 
     texts = [i.page_content for i in texts]
 
-    return texts[0]
+    doc_emb = st_model.encode(texts)
+    doc_emb = doc_emb.tolist()
+
+    ids = [str(uuid.uuid1()) for _ in doc_emb]
+
+    client = chromadb.Client()
+    # Create collection. get_collection, get_or_create_collection, delete_collection also available!
+    collection = client.create_collection("test_db") 
+    
+    collection.add(
+        embeddings=doc_emb,
+        documents=texts,
+        ids=ids,
+        metadata = ["Page": 1, "Section": "diagnosis/prognosis"]
+    )
+
+    return run_query("how to reduce waste?")
+
+
+from transformers import T5ForConditionalGeneration, AutoTokenizer
+import torch
+
+
+model_name = 'google/flan-t5-base'
+
+model = T5ForConditionalGeneration.from_pretrained(model_name, device_map='auto', offload_folder="offload")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+import uuid
+from sentence_transformers import SentenceTransformer
+
+ST_name = 'sentence-transformers/sentence-t5-base'
+
+st_model = SentenceTransformer(ST_name)
+
+
 
 iface = gr.Interface(
     fn=upload_pdf,