add probability calculation

rag_app/rag_2.py

@@ -36,14 +36,14 @@ def completion_to_prompt(completion):
 llm = LlamaCPP(
     model_path="models/Llama-3.2-1B-Instruct-Q4_K_M.gguf",
     temperature=0.1,
-    max_new_tokens=256,
+    max_new_tokens=128,
     context_window=16384,
     model_kwargs={"n_gpu_layers":-1, 'logits_all': False},
     messages_to_prompt=messages_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)
+            n_gpu_layers=-1, n_ctx=8000, logits_all=True)
 
 
 embedding_model = HuggingFaceEmbedding(
@@ -92,24 +92,16 @@ def is_relevant(query, index, threshold=0.7):
 
 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:]:
-        print("evaluating tokens for calculating log probs")
-        llm.eval(eval_tokens)
+    llm.eval(input_tokens)
+    probs = llm.logits_to_logprobs(llm.eval_logits)
+    total_log_prob = 0.0
+    for i, token in enumerate(input_tokens):
+        token_log_prob = probs[i, token]
+        total_log_prob += token_log_prob
         
-        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):
@@ -142,7 +134,7 @@ def answer_question(query):
         retriever=retriever,
         node_postprocessors=[reranker],
     )
-    response = keyword_query_engine.query(query)
+    response = keyword_query_engine.query(f"Answer in less than 100 words: \nQuery:{query}")
     response_text = str(response)
     response_prob = get_sequence_probability(llm2, response_text)
     print(f"Output probability: {response_prob}")