Modified the prompt and confidence score

README.md

@@ -1,5 +1,5 @@
 ---
-title: Slm Financial Rag
+title: SLM Financial RAG
 emoji: 🚀
 colorFrom: green
 colorTo: indigo
@@ -10,5 +10,3 @@ pinned: false
 license: unknown
 short_description: SLM with RAG for Financial Reports
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -37,7 +37,8 @@ os.makedirs("data", exist_ok=True)
 # SLM: Microsoft PHI-2 model is loaded
 # It does have higher memory and compute requirements compared to TinyLlama and Falcon
 # But it gives the best results among the three
-DEVICE = "cpu"  # or cuda
+# DEVICE = "cpu"  # or cuda
+DEVICE = "cuda"  # or cuda
 # MODEL_NAME = "TinyLlama/TinyLlama_v1.1"
 # MODEL_NAME = "tiiuae/falcon-rw-1b"
 MODEL_NAME = "microsoft/phi-2"
@@ -54,7 +55,7 @@ if tokenizer.pad_token is None:
 # Since the model is to be hosted on a cpu instance, we use float32
 # For GPU, we can use float16 or bfloat16
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, torch_dtype=torch.float32, trust_remote_code=True
+    MODEL_NAME, torch_dtype=torch.bfloat16, trust_remote_code=True
 ).to(DEVICE)
 model.eval()
 # model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
@@ -233,6 +234,13 @@ def process_files(files, chunk_size=512):
         pickle.dump(bm25_data, f)
     return "Files processed successfully! You can now query."
 
+def contains_financial_entities(query):
+    """Check if the query has financial entities"""
+    doc = nlp(query)
+    for ent in doc.ents:
+        if ent.label_ in FINANCIAL_ENTITY_LABELS:
+            return True
+    return False
 
 # Input guardrail implementation
 # Regex is used to filter queries related to sensitive topics
@@ -242,7 +250,7 @@ def process_files(files, chunk_size=512):
 def is_query_allowed(query):
     """Checks if the query violates security or confidentiality rules"""
     for pattern in restricted_patterns:
-        if re.search(pattern, query, re.IGNORECASE):
+        if re.search(pattern, query.lower(), re.IGNORECASE):
             return False, "This query requests sensitive or confidential information."
     doc = nlp(query)
     for ent in doc.ents:
@@ -309,45 +317,65 @@ def hybrid_retrieve(query, chunk_texts, index, bm25, top_k=5, lambda_faiss=0.7):
     return final_results
 
 
+def compute_entropy(logits):
+    """Compute entropy from logits."""
+    probs = torch.softmax(logits, dim=-1)
+    log_probs = torch.log(probs + 1e-9)
+    entropy = -(probs * log_probs).sum(dim=-1)
+    return entropy.mean().item()
+
+
 # A confidence score is computed using FAISS and BM25 ranking
-# FAISS: The similarity score between the query (with response) and the retrieved chunks are normalized
-# BM25: The BM25 scores for the query is normalized
-# Both the scores are aggregated using a weighted sum (lambda FAISS) and normalized
-def compute_confidence_score(query, retrieved_chunks, bm25, lambda_faiss):
-    """Calculates a confidence score using FAISS and BM25 rankings."""
+# FAISS: The similarity score between the response and the retrieved chunks are normalized
+# BM25: The BM25 scores for the query and response combined tokens is normalized
+# The mean of top token probability mean and 1-entropy score is the model_conf_signal
+# FAISS, BM25 and the model_conf_signal are combined using a weighted sum
+def compute_response_confidence(
+    query,
+    response,
+    retrieved_chunks,
+    bm25,
+    model_conf_signal=0.5,
+    lambda_faiss=0.4,
+    lambda_conf=0.4,
+    lambda_bm25=1.8,
+):
+    """Calculates a confidence score using FAISS, BM25, top token probabilites and entropy score"""
     if not retrieved_chunks:
-        return 0
-    query_embedding = embed_model.encode(query, normalize_embeddings=True)
-    response_embedding = embed_model.encode(
+        return 0.0
+    # Compute FAISS similarity
+    retrieved_embedding = embed_model.encode(
         " ".join(retrieved_chunks), normalize_embeddings=True
     )
-    # FAISS Similarity
-    faiss_score = np.dot(query_embedding, response_embedding)
+    response_embedding = embed_model.encode(response, normalize_embeddings=True)
+    faiss_score = np.dot(retrieved_embedding, response_embedding)
+    # Normalize the FAISS score
     normalized_faiss = (faiss_score + 1) / 2
-    # BM25 Ranking
-    tokenized_query = query.lower().split()
-    bm25_scores = bm25.get_scores(tokenized_query)
+    # Compute BM25 for combined query + response
+    tokenized_combined = (query + " " + response).lower().split()
+    bm25_scores = bm25.get_scores(tokenized_combined)
+    # Normalize the BM25 score
     if bm25_scores.size > 0:
-        min_bm25 = (
-            np.min(bm25_scores) if np.min(bm25_scores) != np.max(bm25_scores) else 0
-        )
-        max_bm25 = (
-            np.max(bm25_scores) if np.min(bm25_scores) != np.max(bm25_scores) else 1
-        )
-        bm25_score = (
-            np.mean([bm25_scores[idx] for idx in range(len(retrieved_chunks))])
-            if len(retrieved_chunks) > 0
+        bm25_score = np.mean(bm25_scores)
+        min_bm25, max_bm25 = np.min(bm25_scores), np.max(bm25_scores)
+        normalized_bm25 = (
+            (bm25_score - min_bm25) / (max_bm25 - min_bm25 + 1e-6)
+            if min_bm25 != max_bm25
             else 0
         )
-        normalized_bm25 = (bm25_score - min_bm25) / (max_bm25 - min_bm25)
         normalized_bm25 = max(0, min(1, normalized_bm25))
     else:
-        normalized_bm25 = 0
-    # Final Confidence Score (use Lambda FAISS value for weighted sum)
-    confidence_score = round(
-        (normalized_faiss * lambda_faiss + normalized_bm25 * (1 - lambda_faiss)), 2
+        normalized_bm25 = 0.0
+    logger.info(
+        f"Faiss score: {normalized_faiss}, bm25: {normalized_bm25}, "
+        f"Mean Top Token + Entropy Avg: {model_conf_signal}"
+    )
+    confidence_score = (
+        lambda_faiss * normalized_faiss
+        + model_conf_signal * lambda_conf
+        + lambda_bm25 * normalized_bm25
     )
-    return confidence_score
+    return round(min(100, max(0, confidence_score.item() * 100)), 2)
 
 
 # UI handle for query model button
@@ -405,40 +433,62 @@ def query_model(
         else:
             break
     prompt = (
-        f"Based on the following information:\n\n{context}\n\n"
-        "Answer the query in one or two sentences. "
-        "Do not provide follow-ups. "
-        f"Answer the query: {query}"
+        "You are a financial analyst. Answer financial queries concisely using only the numerical data "
+        "explicitly present in the provided financial context:\n\n"
+        f"{context}\n\n"
+        "Strictly use only the given financial data. Do not assume, infer, or generate missing data."
+        " Retain the original format of financial figures exactly as given."
+        " Do not attempt to convert the currency into any other format."
+        " If the requested information is not available in the provided context, respond with "
+        "'No relevant financial data available.'"
+        " Provide exactly one answer in a single sentence."
+        " Do not generate explanations, additional text, or answer multiple queries."
+        f"\nQuery: {query}"
     )
     inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(DEVICE)
     inputs.pop("token_type_ids", None)
     logger.info("Generating output")
     input_len = inputs["input_ids"].shape[-1]
+    logger.info(f"Input len: {input_len}")
     with torch.inference_mode():
         output = model.generate(
             **inputs,
             max_new_tokens=max_new_tokens,
             num_return_sequences=1,
             repetition_penalty=repetition_penalty,
+            output_scores=True,
+            return_dict_in_generate=True,
             pad_token_id=tokenizer.eos_token_id,
         )
-        start_len = 0
-        if use_extraction:
-            start_len = input_len
-        output = output[0][start_len:]
+        sequences = output["sequences"][0][input_len:]
     execution_time = time.perf_counter() - start_time
     logger.info(f"Query processed in {execution_time:.2f} seconds.")
-    response = tokenizer.decode(output, skip_special_tokens=True)
-    confidence_score = compute_confidence_score(
-        query + " " + response, retrieved_chunks, bm25, lambda_faiss
+    log_probs = output["scores"]  # List of logits per generated token
+    token_probs = [torch.softmax(lp, dim=-1) for lp in log_probs]
+    # Extract top token probabilities for each step
+    token_confidences = [tp.max().item() for tp in token_probs]
+    # Compute final confidence score
+    top_token_conf = sum(token_confidences) / len(token_confidences)
+    print(f"Token Token Probability Mean: {top_token_conf:.4f}")
+    entropy_score = sum(compute_entropy(lp) for lp in log_probs) / len(log_probs)
+    entropy_conf = 1 - (entropy_score / torch.log(torch.tensor(tokenizer.vocab_size)))
+    print(f"Entropy-based Confidence: {entropy_conf:.4f}")
+    model_conf_signal = (top_token_conf + (1 - entropy_conf)) / 2
+    response = tokenizer.decode(sequences, skip_special_tokens=True)
+    confidence_score = compute_response_confidence(
+        query, response, retrieved_chunks, bm25, model_conf_signal
     )
-    logger.info(f"Confidence: {confidence_score*100}%")
+    logger.info(f"Confidence: {confidence_score}%")
     if confidence_score <= 0.3:
         logger.error(f"The system is unsure about this response.")
         response += "\nThe system is unsure about this response."
+    final_out = ""
+    if not use_extraction:
+        final_out += f"Context: {context}\nQuery: {query}\n"
+    final_out += f"Response: {response}"
     return (
         response,
-        f"Confidence: {confidence_score*100}%\nTime taken: {execution_time:.2f} seconds",
+        f"Confidence: {confidence_score}%\nTime taken: {execution_time:.2f} seconds",
     )
 
 
@@ -464,7 +514,7 @@ with gr.Blocks(title="Financial Statement RAG with LLM") as ui:
             top_k_input = gr.Number(value=15, label="Top K (Default: 15)")
             lambda_faiss_input = gr.Slider(0, 1, value=0.5, label="Lambda FAISS (0-1)")
             repetition_penalty = gr.Slider(
-                1, 2, value=1.0, label="Repetition Penality (1-2)"
+                1, 2, value=1.2, label="Repetition Penality (1-2)"
             )
             max_tokens_input = gr.Number(value=100, label="Max New Tokens")
         use_extraction = gr.Checkbox(label="Retrieve only the answer", value=False)
@@ -485,6 +535,7 @@ with gr.Blocks(title="Financial Statement RAG with LLM") as ui:
         ],
         outputs=[query_output, time_output],
     )
+
 # Application entry point
 if __name__ == "__main__":
     logger.info("Starting Gradio server...")

data_filters.py

@@ -1,16 +1,16 @@
 """Sensitive data filters"""
 
 restricted_patterns = [
-    r"\b(?:CFO|CEO|CTO|executive|director|manager|employee|staff|worker)\b.*\b(?:salary|compensation|bonus|pay|income)\b",
-    r"\b(?:salary|compensation|bonus|pay|income)\b.*\b(?:CFO|CEO|CTO|executive|director|manager|employee|staff|worker)\b",
+    r"\b(?:cfo|ceo|cto|executive|director|manager|employee|staff|worker)\b.*\b(?:salary|compensation|bonus|pay|income)\b",
+    r"\b(?:salary|compensation|bonus|pay|income)\b.*\b(?:cfo|ceo|cto|executive|director|manager|employee|staff|worker)\b",
     r"\b(?:acquisition|merger|buyout)\b.*\b(?:before|pre-announcement|leak|inside information)\b",
     r"\b(?:before|pre-announcement|leak|inside information)\b.*\b(?:acquisition|merger|buyout)\b",
     r"\b(?:stock price|share price|insider trading|buying shares)\b",
     r"\b(?:internal policy|data breach|security protocol|confidential|classified)\b",
     r"\b(?:password|access credentials|encryption key|secure key)\b",
-    r"\b(?:social security number|SSN|passport number|credit card|bank account|tax ID|TIN|personal details)\b",
-    r"\b(?:employee records|payroll|medical records|HR data|salary data|PII|personally identifiable information)\b",
-    r"\b(?:CFO|CEO|CTO|executive|director|manager|employee|staff|worker)\b.*\b(?:address|work location|home location|residence|personal contact|phone number|email|office location)\b",
+    r"\b(?:social security number|ssn|passport number|credit card|bank account|tax id|tin|personal details)\b",
+    r"\b(?:employee records|payroll|medical records|hr data|salary data|pii|personally identifiable information)\b",
+    r"\b(?:cfo|ceo|cto|executive|director|manager|employee|staff|worker)\b.*\b(?:address|work location|home location|residence|personal contact|phone number|email|office location)\b",
 ]
 
 restricted_topics = {