code/src/llm_classify.py

"""Few-shot stance classification with a chat LLM -- ask it to pick Favor,
Against, or None per (target, tweet), sample it a few times, and turn the
votes into a soft distribution we can blend with the encoders. Few-shot
examples are pulled class-balanced from training, same target preferred.
Uses AUG_BASE_URL / AUG_MODEL like the other LLM scripts.

    python -m src.llm_classify --csv data/track2/dev.csv \\
        --gold data/track2/dev.csv --train data/track2/train.csv \\
        --shots 6 --n 6 --out_probs outputs/llm/t2_dev.npy
"""
import argparse
import json
import os
import re
import urllib.request
from concurrent.futures import ThreadPoolExecutor

import numpy as np

from src.data import (LABEL2ID, TARGET_DESCRIPTIONS,
                      TARGET_DESCRIPTIONS_RICH, load_split)
from src.scorer import load_gold, score

LABELS = ["Against", "Favor", "None"]
SYSTEM = (
    "أنت مصنف موقف عربي دقيق. المطلوب تحديد موقف كاتب التغريدة تجاه "
    "الهدف المحدد. الموقف واحد من ثلاثة فقط:\n"
    "Favor (مؤيد للهدف)، Against (معارض للهدف)، None (لا موقف واضح).\n"
    "أجب بكلمة واحدة فقط من: Favor أو Against أو None."
)
SYSTEM_COT = (
    "أنت خبير في تحليل المواقف في التغريدات العربية، وتراعي اللهجات "
    "والسخرية والتهكم. حلّل موقف كاتب التغريدة تجاه الهدف المحدد على "
    "خطوات موجزة:\n"
    "1) المعنى الحرفي للتغريدة.\n"
    "2) هل فيها سخرية أو تهكم يقلب المعنى الظاهر؟\n"
    "3) الموقف الحقيقي للكاتب تجاه الهدف.\n"
    "ثم في آخر سطر اكتب بالضبط: الموقف: يليه إحدى الكلمات "
    "Favor أو Against أو None."
)
# For a model with a native thinking mode served with --reasoning-parser:
# the chain-of-thought lands in message.reasoning_content and only the final
# answer in message.content. We just ask for a one-word final answer and let
# the model reason internally about dialect / sarcasm / implicit stance.
SYSTEM_THINK = (
    "أنت خبير في تحليل المواقف في التغريدات العربية، تراعي اللهجات الخليجية "
    "والمصرية والشامية والسخرية والتهكم والموقف الضمني غير المباشر. "
    "فكّر ملياً في المعنى الحقيقي لموقف كاتب التغريدة تجاه الهدف المحدد "
    "(انتبه للسخرية التي تقلب الظاهر)، ثم أجب بكلمة واحدة فقط من: "
    "Favor أو Against أو None."
)
_SYSTEM = {"direct": SYSTEM, "cot": SYSTEM_COT, "think": SYSTEM_THINK}


def build_pool(train_df):
    pool = {}
    for tgt, g in train_df.groupby("target"):
        pool[tgt] = {lb: g[g["stance"] == lb]["text"].tolist()
                     for lb in LABELS}
    allc = {lb: train_df[train_df["stance"] == lb]["text"].tolist()
            for lb in LABELS}
    return pool, allc


def pick_shots(target, pool, allc, k, rng):
    per = pool.get(target)
    shots = []
    for i in range(k):
        lb = LABELS[i % 3]
        src = (per[lb] if per and per[lb] else allc[lb])
        if src:
            shots.append((rng.choice(src), lb))
    rng.shuffle(shots)
    return shots


def messages(target, tweet, shots, mode="direct", desc=None):
    system = _SYSTEM.get(mode, SYSTEM)
    tgt = f"{target} ({desc})" if desc else target
    msgs = [{"role": "system", "content": system}]
    for text, lb in shots:
        msgs.append({"role": "user",
                     "content": f"الهدف: {tgt}\nالتغريدة: {text}"})
        ans = f"الموقف: {lb}" if mode == "cot" else lb
        msgs.append({"role": "assistant", "content": ans})
    msgs.append({"role": "user",
                 "content": f"الهدف: {tgt}\nالتغريدة: {tweet}"})
    return msgs


def parse_label(text):
    t = text.strip().lower()
    if "against" in t or "معارض" in t:
        return "Against"
    if "favor" in t or "مؤيد" in t or "مؤيّد" in t:
        return "Favor"
    if re.search(r"\bnone\b", t) or "محايد" in t or "لا موقف" in t:
        return "None"
    return None


def parse_final_label(text):
    """For CoT output: take the label after the last stance marker."""
    tail = text
    for marker in ("الموقف:", "answer:", "stance:"):
        idx = text.lower().rfind(marker)
        if idx != -1:
            tail = text[idx + len(marker):]
            break
    return parse_label(tail)


def call(base_url, model, msgs, n, temperature, max_tokens=8, timeout=180,
         extra_body=None):
    payload = {
        "model": model, "n": n, "temperature": temperature,
        "top_p": 0.95, "max_tokens": max_tokens, "messages": msgs,
    }
    if extra_body:
        payload.update(extra_body)
    body = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        base_url.rstrip("/") + "/chat/completions", data=body,
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        data = json.load(r)
    # With a reasoning parser the answer is in `content`; if the model spent
    # all tokens thinking, fall back to `reasoning_content` so the label is
    # still recoverable.
    outs = []
    for c in data["choices"]:
        m = c["message"]
        outs.append(m.get("content") or m.get("reasoning_content") or "")
    return outs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", required=True)
    ap.add_argument("--train", required=True)
    ap.add_argument("--gold", default=None)
    ap.add_argument("--out_probs", required=True)
    ap.add_argument("--shots", type=int, default=6)
    ap.add_argument("--n", type=int, default=6)
    ap.add_argument("--sets", type=int, default=1,
                    help="ensemble over this many distinct retrieved shot "
                         "sets per query (requires --shot_pool_m)")
    ap.add_argument("--shot_pool_m", type=int, default=0,
                    help="sample the shots from the top-M per class")
    ap.add_argument("--mode", choices=["direct", "cot", "think"],
                    default="direct")
    ap.add_argument("--think_max_tokens", type=int, default=1024,
                    help="token budget for native-thinking mode (reasoning "
                         "+ final one-word answer)")
    ap.add_argument("--enable_thinking", choices=["unset", "true", "false"],
                    default="unset",
                    help="send chat_template_kwargs.enable_thinking; set "
                         "'false' for models that default to thinking-on "
                         "(e.g. Qwen3.6) so direct answers aren't truncated")
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--concurrency", type=int, default=24)
    ap.add_argument("--retrieve", action="store_true")
    ap.add_argument("--rerank", action="store_true",
                    help="rerank retrieved shots with a CrossEncoder")
    ap.add_argument("--rerank_model",
                    default="NAMAA-Space/GATE-Reranker-V1")
    ap.add_argument("--use_desc", action="store_true",
                    help="append a short target description to the prompt")
    ap.add_argument("--rich_desc", action="store_true",
                    help="use the richer target descriptions with --use_desc")
    ap.add_argument("--self_exclude", action="store_true",
                    help="drop a retrieved shot whose text equals the query "
                         "(leave-one-out for transductive pseudo-label pools)")
    ap.add_argument("--embed_model", default="UBC-NLP/MARBERTv2")
    ap.add_argument("--embed_url", default=None,
                    help="OpenAI-compatible /v1 embeddings endpoint; when "
                         "set, retrieval uses it instead of --embed_model")
    ap.add_argument("--embed_api_model", default=None,
                    help="model id for --embed_url")
    ap.add_argument("--embed_instruction", default=None,
                    help="instruction prepended to queries (Qwen3-Embedding)")
    ap.add_argument("--base_url", default=os.environ.get("AUG_BASE_URL", ""))
    ap.add_argument("--model", default=os.environ.get("AUG_MODEL", ""))
    args = ap.parse_args()
    if not args.base_url or not args.model:
        raise SystemExit("set --base_url/--model or AUG_BASE_URL/AUG_MODEL")

    test = load_split(args.csv, "preserve", has_labels=False)
    train = load_split(args.train, "preserve", has_labels=True)
    pool, allc = build_pool(train)
    rng = np.random.default_rng(0)
    seeds = rng.integers(0, 1_000_000, size=len(test))

    retriever, q_emb, reranker = None, None, None
    if args.retrieve:
        import torch
        from src.retrieve import Reranker, Retriever
        device = "cuda" if torch.cuda.is_available() else "cpu"
        retriever = Retriever(train, args.embed_model, device,
                              embed_url=args.embed_url,
                              embed_api_model=args.embed_api_model,
                              instruction=args.embed_instruction)
        q_emb = retriever.embed_queries(
            test["text"].tolist(), args.embed_model, device
        )
        if args.rerank:
            reranker = Reranker(args.rerank_model, device)

    def work(i):
        row = test.iloc[i]
        if args.mode == "think":
            max_tokens, parser = args.think_max_tokens, parse_final_label
        elif args.mode == "cot":
            max_tokens, parser = args.think_max_tokens, parse_final_label
        else:
            max_tokens, parser = 8, parse_label
        # Some models (e.g. Qwen3.6) default to thinking-ON, which starves a
        # short direct answer. Pass enable_thinking explicitly to control it;
        # reasoning (when on) lands in `content` for this server, so cot/think
        # parse the final label from there.
        _tf = {"true": True, "false": False}.get(args.enable_thinking)
        extra_body = ({"chat_template_kwargs": {"enable_thinking": _tf}}
                      if _tf is not None else None)
        desc_map = (TARGET_DESCRIPTIONS_RICH if args.rich_desc
                    else TARGET_DESCRIPTIONS)
        desc = desc_map.get(row["target"]) if args.use_desc else None
        votes = np.zeros(3)
        for e in range(args.sets):
            r = np.random.default_rng(seeds[i] + e)
            if retriever is not None:
                shots = retriever.balanced_shots(
                    q_emb[i], args.shots,
                    query_text=row["text"], reranker=reranker,
                    exclude_text=row["text"] if args.self_exclude else None,
                    sample_m=args.shot_pool_m,
                    rng=r if args.shot_pool_m else None,
                )
            else:
                shots = pick_shots(row["target"], pool, allc, args.shots, r)
            try:
                outs = call(args.base_url, args.model,
                            messages(row["target"], row["text"], shots,
                                     args.mode, desc),
                            args.n, args.temperature, max_tokens,
                            extra_body=extra_body)
            except Exception:
                outs = []
            for o in outs:
                lb = parser(o)
                if lb:
                    votes[LABEL2ID[lb]] += 1
        if votes.sum() == 0:
            votes[LABEL2ID["None"]] = 1
        return votes / votes.sum()

    probs = np.zeros((len(test), 3))
    with ThreadPoolExecutor(max_workers=args.concurrency) as ex:
        for i, p in enumerate(ex.map(work, range(len(test)))):
            probs[i] = p
            if (i + 1) % 200 == 0:
                print(f"  {i + 1}/{len(test)}")

    os.makedirs(os.path.dirname(os.path.abspath(args.out_probs)),
                exist_ok=True)
    np.save(args.out_probs, probs)
    preds = [LABELS[i] for i in probs.argmax(1)]
    print(f"[write] probs -> {args.out_probs}")
    if args.gold:
        score(load_gold(args.gold), preds)


if __name__ == "__main__":
    main()