compute probabilities

rag_app/rag_2.py

@@ -1,12 +1,13 @@
-import re
 import os
-from llama_cpp import Llama, LlamaGrammar
+import math
+import numpy as np
+from llama_cpp import Llama
 from llama_index.llms.llama_cpp import LlamaCPP
 from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
 from llama_index.retrievers.bm25 import BM25Retriever
 from llama_index.core.retrievers import QueryFusionRetriever
 from llama_index.core.query_engine import RetrieverQueryEngine
-from llama_index.core import StorageContext, load_index_from_storage
+from llama_index.core import StorageContext, load_index_from_storage, QueryBundle
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.core.postprocessor import LLMRerank
 from llama_index.core.node_parser import TokenTextSplitter
@@ -37,9 +38,12 @@ llm = LlamaCPP(
     temperature=0.1,
     max_new_tokens=256,
     context_window=16384,
-    model_kwargs={"n_gpu_layers":-1},
+    model_kwargs={"n_gpu_layers":-1, 'logits_all': True, 'logprobs': True,},
     messages_to_prompt=messages_to_prompt,
-    completion_to_prompt=completion_to_prompt)
+    completion_to_prompt=completion_to_prompt,)
+
+llm2 = Llama(model_path="models/Llama-3.2-1B-Instruct-Q4_K_M.gguf",
+            n_gpu_layers=-1, n_ctx=8000, logits_all=True)
 
 
 embedding_model = HuggingFaceEmbedding(
@@ -86,13 +90,39 @@ def is_relevant(query, index, threshold=0.7):
     return not similarity <= threshold
 
 
+def get_sequence_probability(llm, input_sequence):
+    input_tokens = llm.tokenize(input_sequence.encode("utf-8"))
+    sequence_logits = []
+    sequence_logprobs = []
+    
+    eval_tokens = input_tokens[:1]
+    
+    for token in input_tokens[1:]:
+        llm.eval(eval_tokens)
+        
+        probs = llm.logits_to_logprobs(llm.eval_logits)
+        sequence_logits.append(llm.eval_logits[-1][token])
+        sequence_logprobs.append(probs[-1][token])
+        eval_tokens.append(token)
+    
+    total_log_prob = sum(sequence_logprobs)
+    sequence_probability = math.exp(total_log_prob)
+    return sequence_probability
+    
+    
 def answer_question(query):
+    if is_harmful(query):
+        return "This query has been flagged as unsafe."
+
     print("loading bm25 retriever")
     bm25_retriever = BM25Retriever.from_persist_dir("models/bm25_retriever")
     print("loading saved vector index")
     storage_context = StorageContext.from_defaults(persist_dir="models/precomputed_index")
     index = load_index_from_storage(storage_context)
     
+    if not is_relevant(query, index, 0.2):
+        return "This query doesn't appear relevant to finance."
+    
     retriever = QueryFusionRetriever(
         [
             index.as_retriever(similarity_top_k=5, verbose=True),
@@ -111,12 +141,8 @@ def answer_question(query):
         retriever=retriever,
         node_postprocessors=[reranker],
     )
-    
-    if is_harmful(query):
-        return "This query has been flagged as unsafe."
-    
-    if not is_relevant(query, index, 0.2):
-        return "This query doesn't appear relevant to finance."
-    
     response = keyword_query_engine.query(query)
-    return str(response)
+    response_text = str(response)
+    response_prob = get_sequence_probability(llm2, response_text)
+    print(f"Output probability: {response_prob}")
+    return response_text