Update fine_tuning.py

fine_tuning.py

@@ -160,6 +160,38 @@ def load_and_train(model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0"):
     trainer.train()
     model.eval()
 
+# ---------------- Guardrails ----------------
+BLOCKED_TERMS = ["weather", "cricket", "movie", "song", "football", "holiday",
+                 "travel", "recipe", "music", "game", "sports", "politics", "election"]
+
+FINANCE_DOMAINS = [
+    "financial reporting", "balance sheet", "income statement",
+    "assets and liabilities", "equity", "revenue", "profit and loss",
+    "goodwill impairment", "cash flow", "dividends", "taxation",
+    "investment", "valuation", "capital structure", "ownership interests",
+    "subsidiaries", "shareholders equity", "expenses", "earnings",
+    "debt", "amortization", "depreciation"
+]
+finance_embeds = embed_model.encode(FINANCE_DOMAINS, convert_to_tensor=True)
+
+#--------------------------------------------------------------
+#                  GUARD RAIL
+#--------------------------------------------------------------
+def validate_query(query: str, threshold: float = 0.5) -> bool:
+    q_lower = query.lower()
+    if any(bad in q_lower for bad in BLOCKED_TERMS):
+        print("[Guardrail] Rejected by blocklist.")
+        return False
+    q_emb = embed_model.encode(query, convert_to_tensor=True)
+    sim_scores = util.cos_sim(q_emb, finance_embeds)
+    max_score = float(sim_scores.max())
+    if max_score > threshold:
+        print(f"[Guardrail] Accepted (semantic match {max_score:.2f})")
+        return True
+    else:
+        print(f"[Guardrail] Rejected (low semantic score {max_score:.2f})")
+        return False
+
 # -----------------------------
 # GENERATE ANSWER
 # -----------------------------
@@ -167,6 +199,10 @@ def generate_answer(prompt, max_tokens=200):
     if prompt.strip() == "":
         return "Please enter a prompt!"
 
+    if not validate_query(query):
+        print("Query rejected: Not finance-related.")
+        return "Query rejected: Please ask finance-related questions."
+
     system_prompt = "You are a helpful assistant that provides financial data from MakeMyTrip reports."
     messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}]
     input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)