Update app.py

app.py

@@ -103,6 +103,11 @@ def faiss_search(index, question_vector, k=1):
     distances, indices = index.search(question_vector, k)
     return distances,indices
 
+def create_segment_index(vector):
+    segment_index = faiss.IndexFlatL2(vector.shape[1])
+    segment_index.add(vector)
+    return segment_index
+
 
 def predict_faiss(model, tokenizer, embedding_model, df, question, index):
     t = time.time()
@@ -156,12 +161,18 @@ def predict_test(model, tokenizer, embedding_model, df, question, index):  # sen
     mostSimContext = re.sub(r'\s+', ' ', mostSimContext)
 
     segments = sent_tokenize(mostSimContext, engine="crfcut")
+
+    segment_embeddings = get_embeddings(segments)
+    segment_embeddings = prepare_sentences_vector(segment_embeddings)
+    segment_index = create_segment_index(segment_embeddings)
         
-    segments_index = set_index(get_embeddings(embedding_model,segments))
     _distances,_indices = faiss_search(segments_index, question_vector)
     mostSimSegment = segments[_indices[0][0]]
 
     Answer = model_pipeline(model, tokenizer,question,mostSimSegment)
+
+    if len(answer) <= 2:
+        answer = mostSimSegment
     
     # Find the start and end indices of mostSimSegment within mostSimContext
     start_index = mostSimContext.find(Answer)