Update app.py

app.py

@@ -3,7 +3,8 @@ import streamlit as st
 import PyPDF2
 import matplotlib.pyplot as plt
 from io import BytesIO
-from sentence_transformers import SentenceTransformer
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index import Document
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 import dotenv
@@ -17,8 +18,8 @@ dotenv.load_dotenv()
 API_URL = "https://api-inference.huggingface.co/models/sarvamai/sarvam-2b-v0.5"
 headers = {"Authorization": f"Bearer {os.getenv('HUGGINGFACE_API_KEY')}"}
 
-# Initialize SentenceTransformer model
-embed_model = SentenceTransformer('all-MiniLM-L6-v2')
+# Configure embedding model
+embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
 
 def query_huggingface_api(payload):
     response = requests.post(API_URL, headers=headers, json=payload)
@@ -102,9 +103,12 @@ def search_similar_sections(document_text, query, top_k=3):
     # Split the document into sections (you may need to adjust this based on your document structure)
     sections = document_text.split('\n\n')
     
+    # Create Document objects for each section
+    documents = [Document(text=section) for section in sections]
+    
     # Compute embeddings for the query and all sections
-    query_embedding = embed_model.encode([query])[0]
-    section_embeddings = embed_model.encode(sections)
+    query_embedding = embed_model.get_text_embedding(query)
+    section_embeddings = [embed_model.get_text_embedding(doc.text) for doc in documents]
     
     # Compute cosine similarities
     similarities = cosine_similarity([query_embedding], section_embeddings)[0]