addyo07/query-scope-classifier

1#!/usr/bin/env python32"""3Phase 1.3: High-Speed Parallel LLM Semantic Relabeling (Ollama + LMS Dual Engine)4"""5 6import json7import os8import sys9import time10import requests11from concurrent.futures import ThreadPoolExecutor, as_completed12 13INPUT_FILE = "/opt/vox/sandbox/datasets/semantic_raw.jsonl"14OUTPUT_FILE = "/opt/vox/sandbox/datasets/semantic_relabeled.jsonl"15 16OLLAMA_URL = "http://localhost:11434/api/generate"17LMS_URL = "http://localhost:1234/v1/chat/completions"18 19SYSTEM_PROMPT = """Classify the user query into EXACTLY ONE category:20- "User": Personal identity, persona, preferences, user constraints ("My name is Emily", "I am a software engineer", "I prefer async Rust").21- "Domain": Codebases, technical Q&A, active tasks, programming, architecture, bugs ("Fix Tokio deadlock", "How does stage 3 pipeline work?").22- "Temporal": Session recency, context recaps, session continuity ("What did we work on yesterday?", "Summarize last turn").23 24Output JSON ONLY: {"scope": "User" | "Domain" | "Temporal"}"""25 26def classify_query_llm(item, worker_id):27    text = item["text"]28    29    prompt = f"{SYSTEM_PROMPT}\n\nQuery: \"{text}\"\nJSON Output:"30    31    # Alternate between Ollama and LMS depending on worker_id to balance load32    use_lms = (worker_id % 2 == 1)33    34    scope = None35    if use_lms:36        try:37            payload = {38                "model": "llama-3.1-8b-instruct",39                "messages": [{"role": "user", "content": prompt}],40                "temperature": 0.0,41                "response_format": {"type": "json_object"}42            }43            res = requests.post(LMS_URL, json=payload, timeout=10)44            if res.status_code == 200:45                content = res.json()["choices"][0]["message"]["content"]46                parsed = json.loads(content)47                scope = parsed.get("scope")48        except Exception:49            pass50 51    if not scope: # Try Ollama fallback52        try:53            payload = {54                "model": "llama3.1:8b",55                "prompt": prompt,56                "stream": False,57                "options": {"temperature": 0.0}58            }59            res = requests.post(OLLAMA_URL, json=payload, timeout=10)60            if res.status_code == 200:61                resp_text = res.json().get("response", "").strip()62                s = resp_text.find("{")63                e = resp_text.rfind("}")64                if s != -1 and e != -1:65                    parsed = json.loads(resp_text[s:e+1])66                    scope = parsed.get("scope")67        except Exception:68            pass69 70    # Heuristic fast-path fallback if LLM times out71    if not scope or scope not in ["User", "Domain", "Temporal"]:72        text_lower = text.lower()73        if any(w in text_lower for w in ["yesterday", "last session", "previous", "earlier", "recap", "summary", "kal", "pichle"]):74            scope = "Temporal"75        elif any(w in text_lower for w in ["my name", "i am", "i live", "i prefer", "my role", "mera name", "main", "meri"]):76            scope = "User"77        else:78            scope = "Domain" # Primary default79 80    item_copy = dict(item)81    item_copy["scope"] = scope82    return item_copy83 84def main():85    print("=== Phase 1.3: Parallel LLM Semantic Relabeling ===", flush=True)86    87    if not os.path.exists(INPUT_FILE):88        print(f"Error: {INPUT_FILE} missing!", flush=True)89        sys.exit(1)90        91    items = []92    with open(INPUT_FILE, "r", encoding="utf-8") as f:93        for line in f:94            if line.strip():95                items.append(json.loads(line.strip()))96                97    total = len(items)98    print(f"Loaded {total} raw semantic items to relabel.", flush=True)99    100    results = []101    start_time = time.time()102    103    num_workers = 16104    print(f"Starting {num_workers} parallel worker threads across Ollama & LMS...", flush=True)105    106    with ThreadPoolExecutor(max_workers=num_workers) as executor:107        futures = {executor.submit(classify_query_llm, item, idx): idx for idx, item in enumerate(items)}108        109        completed_count = 0110        for future in as_completed(futures):111            res = future.result()112            results.append(res)113            completed_count += 1114            if completed_count % 500 == 0 or completed_count == total:115                elapsed = time.time() - start_time116                rate = completed_count / elapsed117                print(f"  Progress: {completed_count}/{total} ({completed_count/total*100:.1f}%) | {rate:.1f} items/sec", flush=True)118                119    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:120        for item in results:121            f.write(json.dumps(item, ensure_ascii=False) + "\n")122            123    # Distribution tally124    counts = {"User": 0, "Domain": 0, "Temporal": 0}125    for r in results:126        sc = r.get("scope", "Domain")127        counts[sc] = counts.get(sc, 0) + 1128        129    print(f"\n✅ Relabeling complete! Saved {len(results)} items to {OUTPUT_FILE}", flush=True)130    print("Relabeled Scope Distribution:")131    for sc, cnt in counts.items():132        print(f"  - {sc}: {cnt} ({cnt/len(results)*100:.1f}%)", flush=True)133 134if __name__ == "__main__":135    main()136