Create constraint_parser_llm.py

constraint_parser_llm.py

@@ -0,0 +1,192 @@
+# constraint_parser_llm.py
+from __future__ import annotations
+import json, re, os
+from dataclasses import dataclass, asdict
+from typing import Dict, List, Optional
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+# Choose one (Llama3.1 is slightly better at structured JSON)
+MODEL_NAME = os.environ.get("PARSER_MODEL", "meta-llama/Llama-3.1-8B-Instruct")
+# Alternative:
+# MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+
+_tokenizer = None
+_pipe = None
+
+def _lazy_pipe():
+    global _tokenizer, _pipe
+    if _pipe is None:
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            device_map="auto",            # works CPU/GPU
+            trust_remote_code=True,
+        )
+        _pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=_tokenizer,
+            max_new_tokens=256,
+            temperature=0.0,
+            do_sample=False,
+        )
+    return _pipe
+
+DAY_ABBR = ["Sun","Mon","Tue","Wed","Thu","Fri","Sat"]
+DAY_ALIASES = {
+    "sunday":"Sun","sundays":"Sun","sun":"Sun",
+    "monday":"Mon","mondays":"Mon","mon":"Mon",
+    "tuesday":"Tue","tuesdays":"Tue","tue":"Tue",
+    "wednesday":"Wed","wednesdays":"Wed","wed":"Wed",
+    "thursday":"Thu","thursdays":"Thu","thu":"Thu",
+    "friday":"Fri","fridays":"Fri","fri":"Fri",
+    "saturday":"Sat","saturdays":"Sat","sat":"Sat",
+}
+
+@dataclass
+class ParsedConstraints:
+    subject_counts: Dict[str, int]
+    banned_days: List[str]
+    no_before: Optional[str] = None     # "10:00 AM"
+    no_after: Optional[str] = None      # "6:00 PM"
+    keywords: List[str] = None
+    banned_professors: List[str] = None
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self))
+
+SYSTEM_PROMPT = """You convert natural language scheduling requests into STRICT JSON.
+Output only JSON with this schema:
+
+{
+  "subject_counts": {"<SubjectName>": <int>},   // optional
+  "banned_days": ["Sun","Mon","Tue","Wed","Thu","Fri","Sat"], // optional
+  "no_before": "H:MM AM/PM", // optional
+  "no_after": "H:MM AM/PM",  // optional
+  "keywords": ["..."],       // optional
+  "banned_professors": ["..."] // optional
+}
+
+Rules:
+- Never include commentary.
+- Use only "Sun","Mon","Tue","Wed","Thu","Fri","Sat".
+- If the user says 'no weekends', add 'Sat' and 'Sun'.
+- If they say 'mornings'/'afternoons'/'evenings', map to times:
+  * mornings => no_before >= "10:00 AM" (i.e., prefer starts at/after 10am)
+  * afternoons => no_before >= "12:00 PM"
+  * evenings => no_before >= "4:00 PM"
+  * 'no mornings' => no_after <= "12:00 PM"
+- If counts are implied (e.g., "a couple of CS classes"), interpret:
+  couple=2, few=3, several=3.
+- If ambiguous, make reasonable assumptions but keep to schema.
+"""
+
+FEW_SHOTS = [
+    # (user, json)
+    ("make me 4 CS classes, no Sundays",
+     {"subject_counts":{"Computer Science":4},"banned_days":["Sun"]}),
+    ("i want two econ and one psychology after 1pm, avoid friday",
+     {"subject_counts":{"Economics":2,"Psychology":1},"no_before":"1:00 PM","banned_days":["Fri"]}),
+    ("prefer software engineering and algorithms, no class on mon or tue mornings",
+     {"keywords":["software engineering","algorithms"],"banned_days":["Mon","Tue"],"no_after":"12:00 PM"}),
+    ("no weekends, evenings only",
+     {"banned_days":["Sat","Sun"],"no_before":"4:00 PM"}),
+    ("skip wed; a couple math; no classes before 10am",
+     {"subject_counts":{"Mathematics":2},"banned_days":["Wed"],"no_before":"10:00 AM"}),
+]
+
+def _build_prompt(user_text: str) -> str:
+    parts = [SYSTEM_PROMPT, "\nExamples:"]
+    for u, js in FEW_SHOTS:
+        parts.append(f"User: {u}\nJSON: {json.dumps(js)}")
+    parts.append(f"\nUser: {user_text}\nJSON:")
+    return "\n".join(parts)
+
+def _extract_json(text: str) -> Optional[dict]:
+    text = text.strip()
+    # Some models wrap with text; find first/last braces
+    start = text.find("{")
+    end = text.rfind("}")
+    if start != -1 and end != -1 and end > start:
+        try:
+            return json.loads(text[start:end+1])
+        except Exception:
+            pass
+    # Sometimes models produce ```json blocks
+    m = re.search(r"```json\s*(\{.*?\})\s*```", text, re.S)
+    if m:
+        try:
+            return json.loads(m.group(1))
+        except Exception:
+            pass
+    return None
+
+def _normalize_days(days: List[str]) -> List[str]:
+    out = []
+    for d in days or []:
+        k = d.strip()
+        # accept already-correct abbr
+        if k in DAY_ABBR and k not in out:
+            out.append(k); continue
+        # map aliases
+        v = DAY_ALIASES.get(k.lower())
+        if v and v not in out:
+            out.append(v)
+    return out
+
+def _fallback_light(text: str) -> ParsedConstraints:
+    t = (text or "").lower()
+    # crude day catch-all
+    found = set()
+    for k,v in DAY_ALIASES.items():
+        if re.search(rf"\bno( classes?)? (on )?{re.escape(k)}\b", t) or re.search(rf"\bavoid (on )?{re.escape(k)}\b", t) or re.search(rf"\bno {re.escape(k)}\b", t):
+            found.add(v)
+    # time
+    before = re.search(r"no classes before ([0-9]{1,2})(?::([0-9]{2}))?\s*(am|pm)", t)
+    after  = re.search(r"no classes after ([0-9]{1,2})(?::([0-9]{2}))?\s*(am|pm)", t)
+    def fmt(h,m,ap):
+        mm = m if m else "00"
+        return f"{int(h)}:{mm} {ap.upper()}"
+    no_before = fmt(*before.groups()) if before else None
+    no_after  = fmt(*after.groups()) if after else None
+    return ParsedConstraints(subject_counts={}, banned_days=sorted(found), no_before=no_before, no_after=no_after, keywords=[], banned_professors=[])
+
+def parse_constraints(text: str) -> ParsedConstraints:
+    if not text or not text.strip():
+        return ParsedConstraints(subject_counts={}, banned_days=[], keywords=[], banned_professors=[])
+
+    prompt = _build_prompt(text.strip())
+    out = _lazy_pipe()(prompt)[0]["generated_text"]
+
+    obj = _extract_json(out)
+    if not isinstance(obj, dict):
+        # graceful fallback
+        return _fallback_light(text)
+
+    # sanitize / coerce
+    subject_counts = {}
+    for k,v in (obj.get("subject_counts") or {}).items():
+        try:
+            iv = int(v)
+            if iv > 0:
+                subject_counts[k] = iv
+        except Exception:
+            continue
+
+    banned_days = _normalize_days(obj.get("banned_days") or [])
+    no_before = obj.get("no_before") or None
+    no_after  = obj.get("no_after") or None
+
+    def _list_str(x): return [s for s in (x or []) if isinstance(s, str) and s.strip()]
+    keywords = _list_str(obj.get("keywords"))
+    banned_professors = _list_str(obj.get("banned_professors"))
+
+    return ParsedConstraints(
+        subject_counts=subject_counts,
+        banned_days=banned_days,
+        no_before=no_before,
+        no_after=no_after,
+        keywords=keywords,
+        banned_professors=banned_professors
+    )