File size: 4,509 Bytes
7545355 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | import json
import time
import logging
from pathlib import Path
from typing import List, Dict, Any
from tqdm import tqdm
logger = logging.getLogger(__name__)
class TextToMongoGenerator:
def __init__(self, pipeline, config: Any = None):
self.pipeline = pipeline
self.config = config
def _build_prompt(self, question: str, schema: str = "") -> str:
prompt = "### Task: Generate MongoDB Query (MQL)\n"
if schema:
prompt += f"### Collection Schema / Structure:\n{schema}\n\n"
prompt += f"### Question:\n{question}\n\n### MongoDB Query:\n db."
return prompt
def generate_single(self, question: str, schema: str = "") -> Dict[str, Any]:
"""Generates a single MongoDB query and measures latency."""
prompt = self._build_prompt(question, schema)
start_time = time.perf_counter()
raw_output = self.pipeline.generate(prompt)
end_time = time.perf_counter()
latency_ms = (end_time - start_time) * 1000
# Clean up hallucinations and repetition loops
generated_query = self._clean_mongo_output(raw_output)
return {
"generated_query": generated_query,
"latency_ms": round(latency_ms, 2)
}
@staticmethod
def clean_mongo_output(raw_output: str) -> str:
stripped = raw_output.strip()
if stripped.lower().startswith("db."):
# Model re-emitted the prefix itself; don't duplicate it.
restored = stripped
else:
restored = "db." + stripped
cleaned = (
restored.replace("```javascript", "")
.replace("```json", "")
.replace("```js", "")
.replace("```", "")
.strip()
)
# Stop at common hallucination markers
stop_markers = ["--->", "\n\n", "\nUser:", "\nQuestion:", "```", "###", "Human:"]
for marker in stop_markers:
idx = cleaned.find(marker)
if idx != -1:
cleaned = cleaned[:idx]
# Extract the first valid MongoDB query line
lines = cleaned.split('\n')
for line in lines:
line = line.strip()
if line.startswith("db."):
return line.rstrip(';').strip()
# Fallback: find "db." anywhere in the text
if "db." in cleaned:
idx = cleaned.find("db.")
end_idx = len(cleaned)
for marker in stop_markers + ["\n"]:
pos = cleaned.find(marker, idx)
if pos != -1:
end_idx = min(end_idx, pos)
return cleaned[idx:end_idx].strip().rstrip(';')
return cleaned.strip()
_clean_mongo_output = clean_mongo_output
def generate_batch(self, dataset: List[Dict[str, Any]], output_path: str = None) -> List[Dict[str, Any]]:
results = []
logger.info(f"🚀 Starting batch MongoDB query generation for {len(dataset)} examples...")
for item in tqdm(dataset, desc="Generating MongoDB Queries", unit="query"):
query_id = item.get("query_id", f"sample_{len(results)}")
database_id = item.get("database_id", "unknown")
question = item.get("question", "")
schema = item.get("schema", "")
gold_query = item.get("generated_mongodb_query", item.get("gold_query", ""))
if isinstance(gold_query, list) and len(gold_query) > 0:
gold_query = gold_query[0]
if not question:
continue
gen_result = self.generate_single(question, schema)
result_entry = {
"query_id": query_id,
"database_id": database_id,
"question": question,
"gold_query": gold_query,
"generated_query": gen_result["generated_query"],
"latency_ms": gen_result["latency_ms"],
}
results.append(result_entry)
if output_path:
self.save_results(results, output_path)
return results
def save_results(self, results: List[Dict[str, Any]], output_path: str):
path = Path(output_path)
path.parent.mkdir(parents=True, exist_ok=True)
output_data = {"results": results}
with open(path, "w", encoding="utf-8") as f:
json.dump(output_data, f, indent=4, ensure_ascii=False)
logger.info(f"✅ Successfully saved {len(results)} generated queries to {path}")
|