themis / phase1 /eval /classify_intent.py
vg15o2's picture
Moonley backend (HF Space build)
1d9bd9b
Raw
History Blame Contribute Delete
3.54 kB
"""Round-0 intent classifier (the 'router' as a state-setter). Labels each query AUTHORITY
(wants the leading/landmark case on a doctrine -> apply the authority prior) vs SPECIFIC (a
particular case / fact-pattern / narrow holding -> no prior). DeepSeek, parallel, cached to JSON.
Usage: THEMIS_QFILE=authority_queries.tsv OUT=intent_authority.json python classify_intent.py
Reads .env for DEEPSEEK_API_KEY. Cached by qid so re-runs are free; delete the OUT file to refresh.
"""
import os, json, sys, time, concurrent.futures as cf
import requests
HERE = os.path.dirname(os.path.abspath(__file__))
def _load_env(p):
if os.path.exists(p):
for l in open(p):
l = l.strip()
if l and not l.startswith("#") and "=" in l:
k, v = l.split("=", 1); os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))
_load_env(os.path.join(HERE, "..", "scripts", ".env"))
KEY = os.environ["DEEPSEEK_API_KEY"]
HDR = {"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"}
QFILE = os.environ.get("THEMIS_QFILE", "authority_queries.tsv")
OUT = os.environ.get("OUT", "intent.json")
SYS = ("You route a legal search query for an Indian Supreme Court case-law engine. "
"Decide what the user is after:\n"
"AUTHORITY = they want the leading / landmark / controlling case(s) on a legal PRINCIPLE, "
"doctrine, right, or test (e.g. 'is privacy a fundamental right', 'test for sedition', "
"'doctrine of basic structure').\n"
"SPECIFIC = they want a particular named case, a narrow fact-pattern match, a specific "
"statutory provision's application, or a procedural/factual lookup where the single most "
"authoritative landmark is NOT necessarily the right answer.\n"
"Reply with EXACTLY one word: AUTHORITY or SPECIFIC.")
def classify(text):
for attempt in range(3):
try:
r = requests.post("https://api.deepseek.com/chat/completions", headers=HDR, timeout=40,
json={"model": "deepseek-chat", "temperature": 0, "max_tokens": 4,
"messages": [{"role": "system", "content": SYS}, {"role": "user", "content": text}]})
if r.status_code == 200:
t = r.json()["choices"][0]["message"]["content"].strip().upper()
return "AUTHORITY" if "AUTHORITY" in t else "SPECIFIC"
except Exception:
time.sleep(2 * (attempt + 1))
return "SPECIFIC" # fail-safe: no prior
def main():
rows = []
for l in open(QFILE, encoding="utf-8"):
qid, intent, text = l.rstrip("\n").split("\t", 2); rows.append((qid, text))
cache = json.load(open(OUT)) if os.path.exists(OUT) else {}
todo = [(qid, text) for qid, text in rows if qid not in cache]
print(f"{len(rows)} queries, {len(todo)} to classify ({len(cache)} cached)", flush=True)
t0 = time.time()
with cf.ThreadPoolExecutor(max_workers=24) as ex:
futs = {ex.submit(classify, text): qid for qid, text in todo}
done = 0
for f in cf.as_completed(futs):
cache[futs[f]] = f.result(); done += 1
if done % 50 == 0:
json.dump(cache, open(OUT, "w")); print(f" {done}/{len(todo)} {time.time()-t0:.0f}s", flush=True)
json.dump(cache, open(OUT, "w"))
n_auth = sum(1 for v in cache.values() if v == "AUTHORITY")
print(f"done {len(cache)} -> {OUT} | AUTHORITY={n_auth} ({100*n_auth/len(cache):.0f}%) SPECIFIC={len(cache)-n_auth}", flush=True)
if __name__ == "__main__":
main()