Update app.py

app.py

@@ -11,6 +11,8 @@ import uuid
 from sentence_transformers import SentenceTransformer
 import os
 
+globl file_name = ''
+
 model_name = 'google/flan-t5-base'
 model = T5ForConditionalGeneration.from_pretrained(model_name, device_map='auto', offload_folder="offload")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -22,7 +24,7 @@ st_model = SentenceTransformer(ST_name)
 print('sentence read')
 
 
-def get_context(query_text):
+def get_context(query_text, collection):
     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]
@@ -42,8 +44,32 @@ def local_query(query, context):
     return tokenizer.batch_decode(outputs, skip_special_tokens=True)
 
 def run_query(history, query):
+
+
+    loader = PDFMinerLoader(pdf_filename)
+    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]
+
+    client = chromadb.Client()
+    collection = client.create_collection("test_db") 
     
-    context = get_context(query)
+    collection.add(
+        embeddings=doc_emb,
+        documents=texts,
+        ids=ids
+    )
+
+
+    context = get_context(query, collection)
     result = local_query(query, context)
 
     history = history.append(query)
@@ -52,6 +78,7 @@ def run_query(history, query):
 
 def load_document(pdf_filename):
 
+    
     loader = PDFMinerLoader(pdf_filename)
     doc = loader.load()
 
@@ -84,12 +111,10 @@ def upload_pdf(file):
         # Check if the file is not None before accessing its attributes
         if file is not None:
             # Save the uploaded file
-            file_name = file.name
-
-            # file_name = os.path.basename(file_name)
+            file_name = file.name 
 
-            messsage = load_document(file_name)  
-            return messsage
+            # messsage = load_document(file_name)  
+            return 'Successfully uploaded!'
         else:
             return "No file uploaded."