Update app.py

app.py

@@ -7,18 +7,33 @@ import torch
 
 
 
-index = faiss.read_index("index.faiss")
+index = faiss.read_index("index503.faiss")
 
-with open("texts_long.pkl", "rb") as f:
+with open("texts503.pkl", "rb") as f:
     texts = pickle.load(f)
 
+print(texts[500])
 
+# chunking
+chunks = []
+chunk_size = 25
+overlap_int = 2
+for text in texts:
+    text_words = text.split()
+    for i in range(0, len(text_words), chunk_size):
+        chunk_words = text_words[max(i - overlap_int, 0): i + chunk_size]
+        chunk = " ".join(chunk_words)
+        chunks.append(chunk)
+
+print(len(chunks))
+
+# emb model
 model_name = "BAAI/bge-small-en-v1.5"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+emb_model = AutoModel.from_pretrained(model_name)
 
 
-def embed(texts):
+def embed(query):
 
     inputs = tokenizer(
         texts,
@@ -28,7 +43,7 @@ def embed(texts):
     )
 
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = emb_model(**inputs)
 
     token_embeddings = outputs.last_hidden_state
     attention_mask = inputs["attention_mask"]
@@ -49,13 +64,13 @@ def embed(texts):
 
 
 
-def retrieve(query, k=3):
+def retrieve(query, k=5):
 
     q = embed([query])
 
     scores, ids = index.search(q, k)
 
-    return "\n".join(str(i) +": " + texts[i] for i in ids[0])
+    return "\n".join(str(i) +": " + chunks[i] for i in ids[0])
     #return [str(i) +": " + texts[i] + "\n" for i in ids[0]]
 
 
@@ -67,7 +82,7 @@ with gr.Blocks() as demo:
     )
 
     output = gr.Textbox(
-        label="Top3 most relevant news:"
+        label="Top5 most relevant news:"
     )
 
     topic_input.submit(