Update app.py

app.py

@@ -1,44 +1,48 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModel, T5Tokenizer, T5ForConditionalGeneration
+from sentence_transformers import SentenceTransformer
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 import faiss
 import numpy as np
 
-# Load model and tokenizer for sentence transformers
-tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+# Load SentenceTransformer model
+model = SentenceTransformer('all-MiniLM-L6-v2')
 
-# Prepare dataset (Wikipedia dataset can be used)
+# Prepare dataset (Wikipedia dataset or any other dataset can be used)
 corpus = ["Article text 1", "Article text 2", "Article text 3"]
 
-# Tokenize and encode
-encoded_texts = [model(**tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)).last_hidden_state.mean(1).detach().numpy() for text in corpus]
+# Encode the corpus using the sentence-transformers model
+encoded_texts = model.encode(corpus, convert_to_numpy=True)
 
 # Create FAISS index
-dimension = encoded_texts[0].shape[1]
+dimension = encoded_texts.shape[1]
 index = faiss.IndexFlatL2(dimension)
-index.add(np.vstack(encoded_texts))
+index.add(encoded_texts)
 
+# Function to retrieve top-k relevant documents from the corpus
 def retrieve(query, k=5):
-    query_vector = model(**tokenizer(query, return_tensors='pt', truncation=True, max_length=512)).last_hidden_state.mean(1).detach().numpy()
+    query_vector = model.encode([query], convert_to_numpy=True)
     distances, indices = index.search(query_vector, k)
     return [corpus[i] for i in indices[0]]
 
+# Function to generate a human-like response using the FLAN-T5 model
 def generate_response(query):
     retrieved_docs = retrieve(query)
     context = " ".join(retrieved_docs)
     
-    # Use the retrieved context to generate a humanized response
+    # Load the FLAN-T5 model and tokenizer
     flan_t5_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
     flan_t5_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
 
+    # Format the input for the model
     input_text = f"Generate a human-like response: {query}. Context: {context}"
     input_ids = flan_t5_tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512).input_ids
 
-    # Generate text with length constraint
+    # Generate text response with a length constraint
     generated_ids = flan_t5_model.generate(input_ids, max_length=1500)
     response = flan_t5_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
     return response
 
+# Function to trim the generated text to a word limit
 def trim_to_word_limit(text, word_limit=1500):
     words = text.split()
     if len(words) > word_limit:
@@ -60,4 +64,4 @@ if st.button("Generate"):
         st.write(response)
 
 # Additional info or about section
-st.write("This app uses FAISS, sentence-transformers, and FLAN-T5 to generate contextually relevant human-like responses.")
+st.write("This app uses FAISS, SentenceTransformers, and FLAN-T5 to generate contextually relevant human-like responses.")