Update app.py

app.py

@@ -10,7 +10,8 @@ import io
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_core.documents import Document as LCDocument
-from langchain.embeddings import HuggingFaceEmbeddings
+from langchain_core.embeddings import Embeddings
+from sentence_transformers import SentenceTransformer
 import time
 
 # === Load summarization model ===
@@ -20,9 +21,20 @@ model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
 # === Load QA pipeline ===
 qa_pipeline = pipeline("question-answering", model="facebook/bart-large-cnn", tokenizer=tokenizer)
 
-# === Setup BGE Embedding model ===
-embedding_model_name = "BAAI/bge-large-en-v1.5"
-embedding_function = HuggingFaceEmbeddings(model_name=embedding_model_name)
+# === Load SentenceTransformer embedding model ===
+embedding_model = SentenceTransformer("BAAI/bge-small-en-v1.5")
+
+class CustomSentenceTransformer(Embeddings):
+    def __init__(self, model):
+        self.model = model
+
+    def embed_documents(self, texts):
+        return self.model.encode(texts, show_progress_bar=False).tolist()
+
+    def embed_query(self, text):
+        return self.model.encode(text, show_progress_bar=False).tolist()
+
+embedding_function = CustomSentenceTransformer(embedding_model)
 
 # === Data models ===
 class KeyPoint(BaseModel):
@@ -82,7 +94,7 @@ def create_word_report(analysis):
 # === Streamlit UI ===
 st.set_page_config(page_title="Chat With PDF (BART + BGE)", page_icon="📄")
 st.title("📄 Chat With PDF")
-st.caption("Summarize and Chat with Documents using facebook/bart-large-cnn + BGE Embeddings")
+st.caption("Summarize and Chat with Documents using facebook/bart-large-cnn + BGE Small Embeddings")
 
 for key in ["current_file", "pdf_summary", "analysis_time", "pdf_report", "word_report", "vectorstore", "messages"]:
     if key not in st.session_state:
@@ -154,6 +166,6 @@ if st.session_state.vectorstore is not None:
 
 if st.session_state.analysis_time is not None:
     st.markdown(
-        f'<div style="text-align:center; margin-top:2rem; color:gray;">Analysis Time: {st.session_state.analysis_time:.1f}s | Embedding: BGE Large v1.5</div>',
+        f'<div style="text-align:center; margin-top:2rem; color:gray;">Analysis Time: {st.session_state.analysis_time:.1f}s | Embedding: BGE Small v1.5</div>',
         unsafe_allow_html=True
     )