Update app.py

app.py

@@ -1,11 +1,9 @@
 import pandas as pd
-import PyPDF2  # For PDF extraction
+import PyPDF2
 import spacy
-from langchain.chains import ConversationalRetrievalChain
-from langchain.llms import OpenAI
-from langchain.vectorstores import FAISS
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sentence_transformers import SentenceTransformer, util
 import torch
-from transformers import AutoTokenizer, AutoModel
 import gradio as gr
 
 # Load and preprocess PDF text
@@ -27,14 +25,10 @@ df = pd.DataFrame({'text': [pdf_text]})
 # Load the custom embedding model
 class CustomEmbeddingModel:
     def __init__(self, model_name):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModel.from_pretrained(model_name)
+        self.model = SentenceTransformer(model_name)
 
     def embed_text(self, text):
-        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        with torch.no_grad():
-            embeddings = self.model(**inputs).last_hidden_state.mean(dim=1)
-        return embeddings[0].numpy()
+        return self.model.encode(text, convert_to_tensor=True)
 
 embedding_model = CustomEmbeddingModel('distilbert-base-uncased')  # Replace with your model name
 
@@ -43,27 +37,24 @@ nlp = spacy.load("en_core_web_sm")
 
 def preprocess_text(text):
     doc = nlp(text)
-    tokens = [token.lemma_.lower() for token in doc if token.text.lower() not in stopwords.words('english') and token.is_alpha]
+    tokens = [token.lemma_.lower() for token in doc if token.is_alpha]
     return ' '.join(tokens)
 
 # Apply preprocessing and embedding
 df['text'] = df['text'].apply(preprocess_text)
 df['text_embeddings'] = df['text'].apply(lambda x: embedding_model.embed_text(x))
 
-# Create FAISS vector store
-documents = df['text'].tolist()
-embeddings = df['text_embeddings'].tolist()
-vector_store = FAISS.from_documents(documents, embeddings)
-
-# Create LangChain model and chain
-llm_model = OpenAI('gpt-3.5-turbo')  # You can replace this with a different LLM if desired
-retriever = vector_store.as_retriever()
-chain = ConversationalRetrievalChain.from_llm(llm_model, retriever=retriever)
+# Create a FAISS index
+index = faiss.IndexFlatL2(768)  # Assuming embeddings are 768-dimensional
+embeddings = torch.stack(df['text_embeddings'].tolist())
+faiss_index = faiss.IndexFlatL2(embeddings.shape[1])
+faiss_index.add(embeddings.numpy())
 
 # Function to generate a response
 def generate_response(prompt):
-    result = chain({"query": prompt})
-    response = result["result"]
+    query_embedding = embedding_model.embed_text(prompt).unsqueeze(0)
+    distances, indices = faiss_index.search(query_embedding.numpy(), k=1)
+    response = df.iloc[indices[0][0]]['text']
     return response
 
 # Gradio interface