| import os
|
| import re
|
| import json
|
| import torch
|
| import pandas as pd
|
| from sentence_transformers import SentenceTransformer, util
|
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
| import os, json, datetime
|
|
|
| def load_training_data(data_folder):
|
| data = []
|
| missing_code = 0
|
| missing_query = 0
|
| for file in os.listdir(data_folder):
|
| if file.endswith(".json"):
|
| path = os.path.join(data_folder, file)
|
| with open(path, "r", encoding="utf-8") as f:
|
| try:
|
| entries = json.load(f)
|
| for e in entries:
|
| query_text = e.get("english") or e.get("English") or e.get("query") or e.get("Query")
|
| if not query_text:
|
| missing_query += 1
|
| continue
|
|
|
| code_text = e.get("pandas_code") or e.get("Pandas_Code") or e.get("code")
|
| if not code_text:
|
| missing_code += 1
|
| continue
|
|
|
| data.append({"english": query_text, "pandas_code": code_text})
|
| except Exception as e:
|
| print(f" Skipped {file}: {e}")
|
| print(f" Loaded {len(data)} valid query–code pairs from {data_folder}")
|
| print(f" Skipped {missing_code} missing-code and {missing_query} missing-query entries.")
|
| return data
|
|
|
| class EnhancedRetriever:
|
| def __init__(self, data):
|
| self.model = SentenceTransformer("all-MiniLM-L6-v2")
|
| valid_data = [item for item in data if (item.get("pandas_code") and (item.get("english") or item.get("query")))]
|
| if not valid_data:
|
| raise ValueError("No valid query–code pairs found in dataset!")
|
| self.queries = [item.get("english") or item.get("query") for item in valid_data]
|
| self.codes = [item["pandas_code"] for item in valid_data]
|
| print(f" Using {len(valid_data)} valid items for retrieval.")
|
| print(" Encoding queries for retrieval...")
|
| self.query_embeddings = self.model.encode(self.queries, convert_to_tensor=True)
|
|
|
| def retrieve_best_match(self, user_query, top_k=3):
|
| user_emb = self.model.encode(user_query, convert_to_tensor=True)
|
| similarity = util.pytorch_cos_sim(user_emb, self.query_embeddings)[0]
|
| top_results = torch.topk(similarity, k=top_k)
|
| results = []
|
| for i in range(top_k):
|
| results.append({
|
| "query": self.queries[top_results.indices[i]],
|
| "pandas_code": self.codes[top_results.indices[i]],
|
| "similarity": float(top_results.values[i])
|
| })
|
| return results
|
|
|
|
|
|
|
| class Generator:
|
| def __init__(self, model_dir="model"):
|
|
|
| try:
|
| self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
|
| self.model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
|
|
|
|
|
| print(" Base model loaded successfully.")
|
| except Exception as e:
|
| print(f" Error loading base model: {e}")
|
|
|
| self.tokenizer = None
|
| self.model = None
|
|
|
| def generate(self, query):
|
| print("\n================= GENERATE DEBUG =================")
|
| print("Incoming Query :", query)
|
|
|
| try:
|
| prompt = query
|
| print("Prompt Created:", prompt)
|
|
|
| print("Tokenizer Loaded:", self.tokenizer is not None)
|
| try:
|
| inputs = self.tokenizer(prompt, return_tensors="pt")
|
|
|
| except Exception as e:
|
| print(" TOKENIZER ERROR:", e)
|
| return f"[Tokenizer Error] {str(e)}"
|
|
|
|
|
| print("Model Loaded:", self.model is not None)
|
| try:
|
| outputs = self.model.generate(**inputs, max_length=128)
|
| print("Model Generation Successful")
|
| except Exception as e:
|
| print(" MODEL ERROR:", e)
|
| return f"[Model Error] {str(e)}"
|
|
|
|
|
| try:
|
| decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
| print("Decoding Successful →", decoded)
|
| except Exception as e:
|
| print(" DECODING ERROR:", e)
|
| return f"[Decode Error] {str(e)}"
|
|
|
|
|
| return decoded
|
|
|
| except Exception as e:
|
| print(" UNEXPECTED ERROR:", e)
|
| return f"[Unexpected Error] {str(e)}"
|
|
|
| def normalize_name(name):
|
| if not isinstance(name, str):
|
| return name
|
| return re.sub(r'[^a-z0-9]', '', name.lower())
|
|
|
| def extract_column_names(text):
|
| words = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', text)
|
| stopwords = {
|
| 'show','display','find','get','the','and','or','where','what','how','many','much','list','give','me','all','with','for','bottom','top','average','mean','sum','median','count','minimum','maximum','highest','lowest'
|
| }
|
| cols = [w for w in words if w.lower() not in stopwords and len(w) > 2]
|
| return [normalize_name(c) for c in cols]
|
|
|
| def extract_values(text):
|
| quoted = re.findall(r"'([^']*)'", text)
|
| numbers = re.findall(r'\b\d+\b', text)
|
| return quoted + numbers
|
|
|
| def enhanced_adaptation(user_query, code, original_retrieved_query):
|
| query_columns = extract_column_names(user_query)
|
| original_columns = extract_column_names(original_retrieved_query)
|
| query_values = extract_values(user_query)
|
| original_values = extract_values(original_retrieved_query)
|
| new_code = code
|
| for orig_col, new_col in zip(original_columns, query_columns):
|
| if orig_col and new_col and normalize_name(orig_col) != normalize_name(new_col):
|
| for pattern in [rf"'{orig_col}'", rf'\"{orig_col}\"', rf"\\b{orig_col}\\b"]:
|
| new_code = re.sub(pattern, new_col, new_code, flags=re.IGNORECASE)
|
| for orig_val, new_val in zip(original_values, query_values):
|
| if orig_val and new_val and orig_val != new_val:
|
| new_code = re.sub(rf"'{re.escape(orig_val)}'", f"'{new_val}'", new_code)
|
| new_code = re.sub(rf'\"{re.escape(orig_val)}\"', f'\"{new_val}\"', new_code)
|
| new_code = re.sub(rf"\\b{re.escape(orig_val)}\\b", new_val, new_code)
|
| return new_code
|
|
|
| def adapt_operations_based_on_query(query, code):
|
| q = query.lower()
|
| c = code
|
| if any(word in q for word in ["average", "mean", "avg"]):
|
| c = re.sub(r"\.(sum|min|max|count)\(\)", ".mean()", c)
|
| elif any(word in q for word in ["total", "sum", "add", "together"]):
|
| c = re.sub(r"\.(mean|min|max|count)\(\)", ".sum()", c)
|
| return c
|
|
|
| def select_best_template(retrieved_results, user_query):
|
| best_score = -1
|
| best_result = retrieved_results[0]
|
| for result in retrieved_results:
|
| score = result["similarity"]
|
| if score > best_score:
|
| best_score = score
|
| best_result = result
|
| return best_result
|
|
|
| def post_process_code(code, user_query):
|
| code = re.sub(r'\.\.', '.', code)
|
| return code
|
|
|
| def detect_vague_query(user_query):
|
| q = user_query.lower()
|
| intent_keywords = [
|
| "amount", "sum", "total", "average", "mean", "minimum", "maximum",
|
| "count", "number", "list", "show", "display", "details", "records","how many","what is"
|
| ]
|
| data_keywords = [
|
| "claim", "bill", "policy", "approved", "rejected", "tariff", "package", "department", "provider"
|
| ]
|
| has_intent = any(word in q for word in intent_keywords)
|
| has_data = any(word in q for word in data_keywords)
|
| if has_data and not has_intent:
|
| return True
|
| return False
|
|
|
| class RobustHybridText2Code:
|
| def __init__(self, data_folder="data", model_dir="model/my_finetuned_model"):
|
| self.data = load_training_data(data_folder)
|
| self.retriever = EnhancedRetriever(self.data)
|
| self.generator = MultiModelGenerator(
|
| {
|
| "modelA": "model/my_finetuned_model",
|
| "modelB": "model/second_model_folder"
|
| }
|
| )
|
|
|
|
|
|
|
|
|
|
|
| all_cols = set()
|
| for item in self.data:
|
| code = item.get("pandas_code", "")
|
| cols = re.findall(r"df\[['\"]([^'\"]+)['\"]\]", code)
|
| all_cols.update(cols)
|
| self.col_map = {normalize_name(c): c for c in all_cols}
|
|
|
| def get_code(self, user_query):
|
| retrieved_results = self.retriever.retrieve_best_match(user_query, top_k=3)
|
| best = select_best_template(retrieved_results, user_query)
|
| print(f" [Best Match Similarity: {best['similarity']:.2f}]")
|
| print(f" Original Query: {best['query']}")
|
| if best["similarity"] > 0.95:
|
| print(" High similarity — using code directly.")
|
| code = best["pandas_code"]
|
| elif best["similarity"] >= 0.90:
|
| print(" Moderate similarity — adapting code.")
|
| code = enhanced_adaptation(user_query, best["pandas_code"], best["query"])
|
| else:
|
| print(" Low similarity — generating new code.")
|
|
|
|
|
|
|
|
|
| generated, used_model = self.generator.generate(user_query)
|
| print(f"[INFO] Model used for this query → {used_model}")
|
|
|
| code = generated
|
|
|
| return post_process_code(code, user_query)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| class MultiModelGenerator:
|
|
|
| def __init__(self, model_dirs):
|
|
|
| self.model_dirs = {
|
| "modelA": "model/modelA_dir",
|
| "modelB": "model/modelB_dir"
|
| }
|
|
|
| self.models = {}
|
| self.tokenizers = {}
|
|
|
| for name, path in model_dirs.items():
|
| try:
|
| tok = AutoTokenizer.from_pretrained(path)
|
| try:
|
| model = AutoModelForSeq2SeqLM.from_pretrained(path)
|
| except:
|
| from transformers import T5ForConditionalGeneration
|
| model = T5ForConditionalGeneration.from_pretrained(path)
|
|
|
| self.models[name] = model
|
| self.tokenizers[name] = tok
|
| print(f"[MultiModel] Loaded model: {name} from {path}")
|
|
|
| except Exception as e:
|
| print(f"[MultiModel] Failed loading {name}: {e}")
|
|
|
| self.router = SentenceTransformer("all-MiniLM-L6-v2")
|
|
|
|
|
| def pick_best_model(self, query):
|
| query_l = query.lower()
|
|
|
|
|
| trigger_keywords = [
|
| "settelement age",
|
| "outstanding age",
|
| "top",
|
| "claim age",
|
| "rank"
|
| ]
|
|
|
| for kw in trigger_keywords:
|
| if kw in query_l:
|
| print(f"[MultiModel] Keyword '{kw}' detected → Using Model 2")
|
| return "modelB"
|
|
|
|
|
| print("[MultiModel] Default routing → Using Model 1")
|
| return "modelA"
|
|
|
|
|
| def generate(self, query):
|
| print("\n============= MULTI-MODEL GENERATE =============")
|
|
|
| chosen = self.pick_best_model(query)
|
| model = self.models[chosen]
|
| tokenizer = self.tokenizers[chosen]
|
|
|
| inputs = tokenizer(query, return_tensors="pt")
|
| out = model.generate(**inputs, max_length=128)
|
| decoded = tokenizer.decode(out[0], skip_special_tokens=True)
|
|
|
| print(f"[MultiModel] Final Answer from {chosen} → {decoded}")
|
| return decoded, chosen
|
|
|
|
|
|
|
| class ClaimBoticsEngine:
|
| def __init__(self):
|
| print("\n=== Initializing Multi-Dataset + Multi-Model Hybrid Engine ===\n")
|
|
|
|
|
| self.datasets = {
|
| "bill": {"folder": "data/bill_data", "csv_path": "data/bill_data/bill_dataset.csv"},
|
| "status": {"folder": "data/status_data", "csv_path": "data/status_data/dataset.csv"},
|
| "history": {"folder": "data/history_data", "csv_path": "data/history_data/history_dataset.csv"}
|
| }
|
|
|
| self.models_by_dataset = {}
|
| for name, meta in self.datasets.items():
|
| print(f"Loading dataset: {name}")
|
|
|
|
|
| data_pairs = load_training_data(meta["folder"])
|
|
|
|
|
| if os.path.exists(meta["csv_path"]):
|
| df = pd.read_csv(meta["csv_path"])
|
| self.datasets[name]["df"] = df
|
| print(f"Loaded {len(df)} rows for dataset: {name}")
|
| else:
|
| self.datasets[name]["df"] = None
|
| print(f"CSV not found: {meta['csv_path']}")
|
|
|
|
|
| retriever = EnhancedRetriever(data_pairs)
|
|
|
|
|
| generator = MultiModelGenerator({
|
| "modelA": "model/my_finetuned_model",
|
| "modelB": "model/my_second_model"
|
| })
|
|
|
|
|
| self.models_by_dataset[name] = {
|
| "retriever": retriever,
|
| "generator": generator,
|
| "training_data": data_pairs,
|
| }
|
|
|
| print("\nEngine fully loaded!\n")
|
|
|
|
|
| def detect_dataset(self, user_query):
|
| q = user_query.lower()
|
|
|
| if any(x in q for x in ["bill", "billing", "invoice"]):
|
| print("Routing → BILL dataset")
|
| return "bill"
|
|
|
| if any(x in q for x in ["history", "older claims", "previous"]):
|
| print("Routing → HISTORY dataset")
|
| return "history"
|
|
|
| print("Routing → STATUS dataset")
|
| return "status"
|
|
|
|
|
| def execute_code(self, code, df):
|
| try:
|
| local_env = {"df": df, "pd": pd}
|
| result = eval(code, {"__builtins__": {}}, local_env)
|
| return result
|
| except Exception as e:
|
| print("Execution error:", e)
|
| return None
|
|
|
|
|
| def process_query(self, user_query):
|
|
|
|
|
| dataset = self.detect_dataset(user_query)
|
| df = self.datasets[dataset]["df"]
|
| retriever = self.models_by_dataset[dataset]["retriever"]
|
| generator = self.models_by_dataset[dataset]["generator"]
|
|
|
|
|
| retrieved = retriever.retrieve_best_match(user_query, top_k=3)
|
| best = select_best_template(retrieved, user_query)
|
|
|
| print(f"Best similarity = {best['similarity']:.4f}")
|
|
|
|
|
| if best["similarity"] > 0.95:
|
| print("High similarity → using template")
|
| code = best["pandas_code"]
|
|
|
|
|
| elif best["similarity"] >= 0.90:
|
| print("Medium similarity → adapting template")
|
| code = enhanced_adaptation(user_query, best["pandas_code"], best["query"])
|
|
|
|
|
| else:
|
| print("Low similarity → generating via LLM (ModelA / ModelB)")
|
| code, used_model = generator.generate(user_query)
|
| print(f"LLM Model Used: {used_model}")
|
|
|
|
|
| code = post_process_code(code, user_query)
|
|
|
|
|
| result = None
|
| if df is not None:
|
| result = self.execute_code(code, df)
|
|
|
|
|
| log_interaction(user_query, code, result, dataset)
|
|
|
| return {
|
| "dataset": dataset,
|
| "code": code,
|
| "result": result
|
| }
|
|
|
|
|
| def log_interaction(query, code, result, dataset):
|
|
|
| os.makedirs("logs", exist_ok=True)
|
| log_path = "logs/interaction_log.jsonl"
|
| entry = {
|
| "timestamp": datetime.datetime.now().isoformat(),
|
| "query": query,
|
| "dataset": dataset,
|
| "code": code,
|
| "result_preview": str(result)[:300]
|
| }
|
| with open(log_path, "a", encoding="utf-8") as f:
|
| f.write(json.dumps(entry) + "\n")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| if __name__ == "__main__":
|
| print("\nInitializing Hybrid Text-to-Pandas Bot...\n")
|
| bot = ClaimBoticsEngine()
|
|
|
| print("Chatbot ready! Type your queries below. Type 'exit' to quit.\n")
|
|
|
| while True:
|
| try:
|
| user_query = input("You: ").strip()
|
| if user_query.lower() in ["exit", "quit"]:
|
| print("Goodbye!")
|
| break
|
| if not user_query:
|
| continue
|
|
|
| code, result = bot.get_code(user_query)
|
|
|
| print("\nBot:")
|
| if code:
|
| print(f" Suggested Pandas Code:\n {code}")
|
| else:
|
| print(" Could not generate code. Please clarify your query.")
|
|
|
| if result is not None:
|
| print(f"\n Execution Result:\n {result}")
|
| print("\n" + "-"*50 + "\n")
|
|
|
| except KeyboardInterrupt:
|
| print("\nInterrupted. Exiting chatbot...")
|
| break
|
| except Exception as e:
|
| print(f"Error: {e}")
|
| continue
|
|
|
|
|
|
|
|
|
|
|
