Update app.py

app.py

@@ -1,11 +1,11 @@
-from dotenv import load_dotenv
-import streamlit as st
-import pickle
-from PyPDF2 import PdfReader
-from transformers import pipeline, AutoTokenizer, AutoModel
 import os
-import torch
+import pickle
 import numpy as np
+from PyPDF2 import PdfReader
+from transformers import pipeline
+from sentence_transformers import SentenceTransformer
+from dotenv import load_dotenv
+import streamlit as st
 
 # Load environment variables from .env file
 load_dotenv()
@@ -15,24 +15,14 @@ def chunk_text(text, chunk_size=1000, chunk_overlap=200):
     chunks = []
     i = 0
     while i < len(text):
-        # Ensure chunk size and overlap are handled properly
         chunks.append(text[i:i + chunk_size])
         i += chunk_size - chunk_overlap
     return chunks
 
-# Function to generate embeddings using transformers
-def generate_embeddings(text_chunks, model_name='sentence-transformers/all-MiniLM-L6-v2'):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModel.from_pretrained(model_name)
-
-    embeddings = []
-    for text in text_chunks:
-        # Tokenize the text and generate embeddings
-        inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
-        with torch.no_grad():
-            outputs = model(**inputs)
-        # Mean pooling on the last hidden state
-        embeddings.append(outputs.last_hidden_state.mean(dim=1).squeeze().numpy())
+# Function to generate embeddings using sentence-transformers
+def generate_embeddings(text_chunks, model_name='all-MiniLM-L6-v2'):
+    model = SentenceTransformer(model_name)
+    embeddings = model.encode(text_chunks, convert_to_tensor=False)
     return embeddings
 
 # Function to find the most relevant chunk based on the cosine similarity
@@ -52,7 +42,6 @@ def main():
 
     if pdf is not None:
         pdf_reader = PdfReader(pdf)
-
         text = ""
         for page in pdf_reader.pages:
             text += page.extract_text()
@@ -89,8 +78,8 @@ def main():
             result = qa_pipeline(question=query, context=best_chunk)
             st.write(result['answer'])
 
-if __name__ == '__main__':
-    main()
+    # Set background image from URL
+    set_bg_from_url("https://www.1access.com/wp-content/uploads/2019/10/GettyImages-1180389186.jpg", opacity=0.5)
 
 def set_bg_from_url(url, opacity=1):
     footer = """
@@ -116,20 +105,5 @@ def set_bg_from_url(url, opacity=1):
     </footer>
     """
     st.markdown(footer, unsafe_allow_html=True)
-
-    # Set background image using HTML and CSS
-    st.markdown(
-        f"""
-        <style>
-            body {{
-                background: url('{url}') no-repeat center center fixed;
-                background-size: cover;
-                opacity: {opacity};
-            }}
-        </style>
-        """,
-        unsafe_allow_html=True
-    )
-
-# Set background image from URL
-set_bg_from_url("https://www.1access.com/wp-content/uploads/2019/10/GettyImages-1180389186.jpg", opacity=0.875)
+    
+    # Set background image using