Update app.py

app.py

@@ -1,5 +1,9 @@
 import os
 import torch
+import uuid
+import shutil
+import numpy as np
+import faiss
 from flask import Flask, jsonify, request
 from flask_cors import CORS
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig
@@ -30,6 +34,55 @@ accelerator = Accelerator()
 highest_label = None 
 loaded_models = {}
  
+# FAISS Index Setup
+UPLOAD_DIR = "./uploads"
+faiss_index_file = os.path.join(UPLOAD_DIR, "faiss_index.bin")
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+document_store = {}
+
+if os.path.exists(faiss_index_file):
+    try:
+        index = faiss.read_index(faiss_index_file)
+        if index.ntotal > 0:
+            print(f"FAISS index loaded with {index.ntotal} vectors.")
+            index.reset()
+            index = faiss.IndexIDMap(faiss.IndexFlatL2(384))
+        else:
+            index = faiss.IndexIDMap(faiss.IndexFlatL2(384))
+    except Exception as e:
+        print(f"Error loading FAISS index: {e}, reinitializing.")
+        index = faiss.IndexIDMap(faiss.IndexFlatL2(384))
+else:
+    index = faiss.IndexIDMap(faiss.IndexFlatL2(384))
+
+# Function to upload document
+def upload_document(file_path, embed_model):
+    doc_id = uuid.uuid4().int % (2**63 - 1)
+    file_location = os.path.join(UPLOAD_DIR, os.path.basename(file_path))
+    shutil.copy(file_path, file_location)
+    
+    with open(file_location, "r", encoding="utf-8") as f:
+        text = f.read()
+    
+    vector = embed_model.encode(text).astype("float32")
+    index.add_with_ids(np.array([vector]), np.array([doc_id], dtype=np.int64))
+    document_store[doc_id] = {"path": file_location, "text": text}
+    
+    faiss.write_index(index, faiss_index_file)
+    print(f"Document uploaded with doc_id: {doc_id}")
+
+@app.route("/upload", methods=["POST"])
+def handle_upload():
+    if "file" not in request.files:
+        return jsonify({"error": "No file provided"}), 400
+    
+    file = request.files["file"]
+    file_path = os.path.join(UPLOAD_DIR, file.filename)
+    file.save(file_path)
+    
+    upload_document(file_path, bertmodel)
+    return jsonify({"message": "File uploaded successfully"})
+
 def get_model_and_tokenizer(model_id: str):
     """
     Load and cache the model and tokenizer for the given model_id.
@@ -108,6 +161,17 @@ def generate_response(user_input, model_id):
        
         # Append chat history
         func_caller = []
+        
+        query_vector = bertmodel.encode(user_input).reshape(1, -1).astype("float32")
+        D, I = index.search(query_vector, 1)
+        
+        retrieved_knowledge = document_store.get(I[0][0], {}).get("text", "No relevant information found.")
+        
+        # Construct the knowledge prompt
+        prompt = f"Use the following knowledge:\n{retrieved_knowledge}"
+        
+        # Add the retrieved knowledge to the prompt
+        func_caller.append({"role": "system", "content": prompt}) 
 
         for msg in chat_history:
             func_caller.append({"role": "user", "content": f"{str(msg[0])}"})