Update app.py

app.py

@@ -8,11 +8,10 @@ tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v
 model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
 
 # Prepare dataset (Wikipedia dataset can be used)
-# Example: [title, text] pairs
 corpus = ["Article text 1", "Article text 2", "Article text 3"]
 
 # Tokenize and encode
-encoded_texts = [model(**tokenizer(text, return_tensors='pt', padding=True)).last_hidden_state.mean(1).detach().numpy() for text in corpus]
+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]
 
 # Create FAISS index
 dimension = encoded_texts[0].shape[1]
@@ -20,7 +19,7 @@ index = faiss.IndexFlatL2(dimension)
 index.add(np.vstack(encoded_texts))
 
 def retrieve(query, k=5):
-    query_vector = model(**tokenizer(query, return_tensors='pt')).last_hidden_state.mean(1).detach().numpy()
+    query_vector = model(**tokenizer(query, return_tensors='pt', truncation=True, max_length=512)).last_hidden_state.mean(1).detach().numpy()
     distances, indices = index.search(query_vector, k)
     return [corpus[i] for i in indices[0]]
 
@@ -33,7 +32,7 @@ def generate_response(query):
     flan_t5_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
 
     input_text = f"Generate a human-like response: {query}. Context: {context}"
-    input_ids = flan_t5_tokenizer(input_text, return_tensors="pt").input_ids
+    input_ids = flan_t5_tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512).input_ids
 
     # Generate text with length constraint
     generated_ids = flan_t5_model.generate(input_ids, max_length=1500)