change reranker

rag_app/rag_2.py

@@ -9,13 +9,39 @@ from llama_index.core.query_engine import RetrieverQueryEngine
 from llama_index.core import StorageContext, load_index_from_storage
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.core.postprocessor import LLMRerank
+from llama_index.core.node_parser import TokenTextSplitter
+from transformers import AutoTokenizer
+from llama_index.core.postprocessor import SentenceTransformerRerank
+
+_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
+
+
+def messages_to_prompt(messages):
+    messages = [{"role": m.role.value, "content": m.content} for m in messages]
+    prompt = _tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    return prompt
+
+
+def completion_to_prompt(completion):
+    messages = [{"role": "user", "content": completion}]
+    prompt = _tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    return prompt
+
 
 llm = LlamaCPP(
     model_path="models/Llama-3.2-1B-Instruct-Q4_K_M.gguf",
     temperature=0.1,
     max_new_tokens=256,
-    context_window=16384
-)
+    context_window=16384,
+    model_kwargs={"n_gpu_layers":-1},
+    messages_to_prompt=messages_to_prompt,
+    completion_to_prompt=completion_to_prompt)
+
+
 embedding_model = HuggingFaceEmbedding(
     model_name="models/all-MiniLM-L6-v2"
 )
@@ -34,11 +60,15 @@ def check_if_exists():
 
 def precompute_index(data_folder='data'):
     documents = SimpleDirectoryReader(data_folder).load_data()
-    index = VectorStoreIndex.from_documents(documents)
+    splitter = TokenTextSplitter(chunk_size=400, chunk_overlap=50)
+    nodes = splitter.get_nodes_from_documents(documents)
+    index = VectorStoreIndex(nodes, verbose=True)
+    # index = VectorStoreIndex.from_documents(documents)
     index.storage_context.persist(persist_dir='models/precomputed_index')
     bm25_retriever = BM25Retriever.from_defaults(
-        nodes=documents,
-        similarity_top_k=5
+        nodes=nodes,
+        similarity_top_k=5,
+        verbose=True
     )
     bm25_retriever.persist("models/bm25_retriever")
 
@@ -56,20 +86,21 @@ def answer_question(query):
     
     retriever = QueryFusionRetriever(
         [
-            index.as_retriever(similarity_top_k=5),
+            index.as_retriever(similarity_top_k=5, verbose=True),
             bm25_retriever,
         ],
         llm=llm,
         num_queries=1,
         similarity_top_k=5,
+        verbose=True
     )
-    reranker = LLMRerank(
-        choice_batch_size=5,
-        top_n=5,
+    reranker = SentenceTransformerRerank(
+        model="cross-encoder/ms-marco-MiniLM-L-2-v2",
+        top_n=5
     )
     keyword_query_engine = RetrieverQueryEngine(
         retriever=retriever,
-        node_postprocessors=[reranker]
+        node_postprocessors=[reranker],
     )
     
     if is_harmful(query):