Update app.py

app.py

@@ -1,6 +1,4 @@
 import os
-import multiprocessing
-import concurrent.futures
 from langchain.document_loaders import TextLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
@@ -10,90 +8,51 @@ import torch
 import numpy as np
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from datetime import datetime
-import json
 import gradio as gr
 import re
 from threading import Thread
 
-from llama_index.core import VectorStoreIndex, Document
-from llama_index.core.tools import QueryEngineTool, ToolMetadata
-from llama_index.agent.openai import OpenAIAgent
-from llama_index.llms.openai import OpenAI
-from llama_index.embeddings.openai import OpenAIEmbedding
-
 class MultiDocumentAgentSystem:
-    def __init__(self, documents_dict, llm, embed_model):
-        self.llm = llm
-        self.embed_model = embed_model
-        self.document_agents = {}
-        self.create_document_agents(documents_dict)
-        self.top_agent = self.create_top_agent()
-
-    def create_document_agents(self, documents_dict):
-        for doc_name, doc_content in documents_dict.items():
-            vector_index = VectorStoreIndex.from_documents([Document(text=doc_content)])
-            summary_index = VectorStoreIndex.from_documents([Document(text=doc_content)])
-            
-            vector_query_engine = vector_index.as_query_engine(similarity_top_k=2)
-            summary_query_engine = summary_index.as_query_engine()
-
-            query_engine_tools = [
-                QueryEngineTool(
-                    query_engine=vector_query_engine,
-                    metadata=ToolMetadata(
-                        name=f"vector_tool_{doc_name}",
-                        description=f"Useful for specific questions about {doc_name}",
-                    ),
-                ),
-                QueryEngineTool(
-                    query_engine=summary_query_engine,
-                    metadata=ToolMetadata(
-                        name=f"summary_tool_{doc_name}",
-                        description=f"Useful for summarizing content about {doc_name}",
-                    ),
-                ),
-            ]
-
-            self.document_agents[doc_name] = OpenAIAgent.from_tools(
-                query_engine_tools,
-                llm=self.llm,
-                verbose=True,
-                system_prompt=f"You are an agent designed to answer queries about {doc_name}.",
-            )
-
-    def create_top_agent(self):
-        all_tools = []
-        for doc_name, agent in self.document_agents.items():
-            doc_tool = QueryEngineTool(
-                query_engine=agent,
-                metadata=ToolMetadata(
-                    name=f"tool_{doc_name}",
-                    description=f"Use this tool for questions about {doc_name}",
-                ),
-            )
-            all_tools.append(doc_tool)
-
-        obj_index = VectorStoreIndex.from_objects(all_tools, embed_model=self.embed_model)
-
-        return OpenAIAgent.from_tools(
-            all_tools,
-            llm=self.llm,
-            verbose=True,
-            system_prompt="You are an agent designed to answer queries about multiple documents.",
-            tool_retriever=obj_index.as_retriever(similarity_top_k=3),
-        )
+    def __init__(self, documents_dict, model, tokenizer, embeddings):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.embeddings = embeddings
+        self.document_vectors = self.create_document_vectors(documents_dict)
+
+    def create_document_vectors(self, documents_dict):
+        document_vectors = {}
+        for doc_name, content in documents_dict.items():
+            vectors = self.embeddings.encode(content, convert_to_tensor=True)
+            document_vectors[doc_name] = vectors
+        return document_vectors
 
     def query(self, user_input):
-        return self.top_agent.chat(user_input)
+        query_vector = self.embeddings.encode(user_input, convert_to_tensor=True)
+        
+        # Find the most similar document
+        most_similar_doc = max(self.document_vectors.items(), 
+                               key=lambda x: torch.cosine_similarity(query_vector, x[1], dim=0))
+        
+        # Generate response using the most similar document as context
+        response = self.generate_response(user_input, most_similar_doc[0], most_similar_doc[1])
+        return response
+
+    def generate_response(self, query, doc_name, doc_vector):
+        prompt = f"Based on the document '{doc_name}', answer the following question: {query}"
+        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.model.device)
+        
+        with torch.no_grad():
+            output = self.model.generate(input_ids, max_length=150, num_return_sequences=1)
+        
+        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return response
 
 class DocumentRetrievalAndGeneration:
     def __init__(self, embedding_model_name, lm_model_id, data_folder):
         self.documents_dict = self.load_documents(data_folder)
         self.embeddings = SentenceTransformer(embedding_model_name)
         self.tokenizer, self.model = self.initialize_llm(lm_model_id)
-        self.llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
-        self.embed_model = OpenAIEmbedding()
-        self.multi_doc_system = MultiDocumentAgentSystem(self.documents_dict, self.llm, self.embed_model)
+        self.multi_doc_system = MultiDocumentAgentSystem(self.documents_dict, self.model, self.tokenizer, self.embeddings)
 
     def load_documents(self, folder_path):
         documents_dict = {}
@@ -102,7 +61,7 @@ class DocumentRetrievalAndGeneration:
                 file_path = os.path.join(folder_path, file_name)
                 with open(file_path, 'r', encoding='utf-8') as file:
                     content = file.read()
-                    documents_dict[file_name[:-4]] = content  # Use filename without .txt as key
+                    documents_dict[file_name[:-4]] = content
         return documents_dict
 
     def initialize_llm(self, model_id):
@@ -132,7 +91,7 @@ class DocumentRetrievalAndGeneration:
                 top_k=20,
                 temperature=0.8,
                 repetition_penalty=1.2,
-                eos_token_id=[128001, 128008, 128009],
+                eos_token_id=self.tokenizer.eos_token_id,
                 streamer=streamer,
             )
             
@@ -157,8 +116,8 @@ class DocumentRetrievalAndGeneration:
         return response, related_queries
 
 if __name__ == "__main__":
-    embedding_model_name = 'flax-sentence-embeddings/all_datasets_v3_MiniLM-L12'
-    lm_model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    embedding_model_name = 'sentence-transformers/all-MiniLM-L6-v2'
+    lm_model_id = "facebook/opt-350m"  # You can change this to a different open-source model
     data_folder = 'sample_embedding_folder2'
 
     doc_retrieval_gen = DocumentRetrievalAndGeneration(embedding_model_name, lm_model_id, data_folder)