Blame for scripts/kgeval_task1_qwen_hosted.py

# -*- coding: utf-8 -*-
"""KnowledgeGraphEval 2026 · Subtask 1 (AdaptNER) — generative NER via a HOSTED QWEN endpoint.

Standalone, local, resume-friendly. Mirrors the AlexandriaX `alexandriax_task1_qwen_hosted.py`
pattern: instead of training, it drives an existing **hosted Qwen inference endpoint**
(OpenAI-compatible, e.g. the vLLM server on the femad/data-intelligence stack) to extract
named entities from the Konooz dev/eval sentences, then writes the official CodaBench
submission (token + 21 B/I/O columns, one concatenated file).

Design
------
* No training / no GPU. Pure inference against an existing hosted model.
* Runs on your PC. Outputs saved next to this file under ``qwen_runs/``.
* Direct OpenAI-compatible ``/v1/chat/completions`` call — no orchestrator, no tracing.
* Resume-safe. Each sentence's raw model output is cached to a local JSONL as it lands;
  re-run and it only queries the sentences still missing.

Endpoint config (env vars — sane defaults)
------------------------------------------
    QWEN_BASE_URL     OpenAI-compatible base URL. Default: http://localhost:8005/v1
                      (reach a remote server via: ssh -N -L 8005:localhost:8005 <host>)
    QWEN_MODEL        Model name served by the endpoint. Default: qwen3-3b
                      (SET THIS to whatever your Qwen 3B is registered as; check
                      with:  curl http://localhost:8005/v1/models )
    QWEN_API_KEY      vLLM ignores it; any non-empty string. Default: "EMPTY".
    QWEN_HOST_HEADER  Proxy vhost when tunnelling (e.g. "100.64.0.2:8005"). Default: none.
    QWEN_NO_THINK     "1" (default) disables qwen3 thinking mode (harmless on Instruct).
    QWEN_CONCURRENCY  Parallel in-flight requests. Default: 8.
    QWEN_TEMPERATURE  Sampling temperature. Default: 0.0 (deterministic).
    QWEN_MAX_TOKENS   Max generated tokens per sentence. Default: 512.
    QWEN_MAX_RETRIES  Transient-error retries. Default: 4.

Data config
-----------
    KONOOZ_SRC        Path to a .zip OR a folder of <Domain>.txt files.
                      Default: auto-find under ../../data/raw/konooz_dev then dev-konooz.zip.

Run
---
    pip install openai tqdm
    QWEN_MODEL=qwen2-2b python kgeval_task1_qwen_hosted.py
"""
from __future__ import annotations

import glob
import io
import json
import os
import re
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed

from tqdm.auto import tqdm

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
ROOT = os.path.dirname(os.path.abspath(__file__))
WORK_DIR = os.path.join(ROOT, "qwen_runs")
SUBMIT_DIR = os.path.join(WORK_DIR, "submissions_qwen")
os.makedirs(SUBMIT_DIR, exist_ok=True)

NAME = os.environ.get("RUN_NAME", "qwen3b_hosted")
CACHE_FILE = os.path.join(SUBMIT_DIR, f"cache_{NAME}.jsonl")

# Official 21 entity types, in submission column order.
TYPE_ORDER = [
    "CARDINAL", "CURR", "DATE", "EVENT", "FAC", "GPE", "LANGUAGE", "LAW", "LOC", "MONEY",
    "NORP", "OCC", "ORDINAL", "ORG", "PERCENT", "PERS", "PRODUCT", "QUANTITY", "TIME", "UNIT", "WEBSITE",
]

# --- hosted qwen endpoint (direct, untraced) ---
QWEN_BASE_URL = os.environ.get("QWEN_BASE_URL") or os.environ.get("LLM_API_BASE_URL") or "http://localhost:8005/v1"
QWEN_MODEL = os.environ.get("QWEN_MODEL") or os.environ.get("LLM_MODEL_NAME") or "qwen3-3b"
QWEN_API_KEY = os.environ.get("QWEN_API_KEY", "EMPTY")
QWEN_HOST_HEADER = os.environ.get("QWEN_HOST_HEADER", "")
QWEN_NO_THINK = os.environ.get("QWEN_NO_THINK", "1") == "1"
CONCURRENCY = int(os.environ.get("QWEN_CONCURRENCY", "8"))
TEMPERATURE = float(os.environ.get("QWEN_TEMPERATURE", "0.0"))
MAX_TOKENS = int(os.environ.get("QWEN_MAX_TOKENS", "512"))
MAX_RETRIES = int(os.environ.get("QWEN_MAX_RETRIES", "4"))

KONOOZ_SRC = os.environ.get("KONOOZ_SRC", "")

# ---------------------------------------------------------------------------
# Data loading (CoNLL: token + N columns; blank line = sentence boundary)
# ---------------------------------------------------------------------------
def resolve_konooz_src() -> str:
    if KONOOZ_SRC:
        return KONOOZ_SRC
    guesses = [
        os.path.join(ROOT, "..", "..", "data", "raw", "konooz_dev"),
        os.path.join(ROOT, "..", "..", "data", "raw", "dev-konooz.zip"),
        os.path.join(ROOT, "..", "..", "data", "raw", "devkonooz.zip"),
    ]
    for g in guesses:
        if os.path.exists(g):
            return g
    raise FileNotFoundError("Set KONOOZ_SRC to the Konooz dev .zip or folder.")


def load_domain_files(src: str):
    if str(src).endswith(".zip"):
        work = os.path.join(WORK_DIR, "konooz_dev")
        os.makedirs(work, exist_ok=True)
        with zipfile.ZipFile(src) as z:
            z.extractall(work)
        src = work
    files = sorted(glob.glob(os.path.join(src, "**", "*.txt"), recursive=True))
    if not files:
        raise FileNotFoundError(f"no .txt files under {src}")
    return files


def read_conll(path: str):
    sents, cur = [], []
    for line in io.open(path, encoding="utf-8"):
        s = line.rstrip("\n")
        if not s.strip():
            if cur:
                sents.append(cur)
                cur = []
            continue
        cur.append(s.split())  # parts[0] = token, parts[1:] = placeholder tags (ignored)
    if cur:
        sents.append(cur)
    return sents


# ---------------------------------------------------------------------------
# Qwen client (direct OpenAI-compatible call)
# ---------------------------------------------------------------------------
from openai import OpenAI

_client = OpenAI(
    base_url=QWEN_BASE_URL,
    api_key=QWEN_API_KEY,
    max_retries=0,
    default_headers={"Host": QWEN_HOST_HEADER} if QWEN_HOST_HEADER else None,
)

_TYPE_LIST_STR = ", ".join(TYPE_ORDER)
_SYSTEM = (
    "You are an expert Arabic Named Entity Recognition system. "
    "Given an Arabic sentence, extract every named entity. "
    "Allowed entity types (use these labels exactly): " + _TYPE_LIST_STR + ". "
    "Entities may be nested (a span can belong to more than one type). "
    "Output ONLY the entities, one per line, as:  <surface text><TAB><TYPE>. "
    "Copy the surface text exactly as it appears in the sentence. "
    "If there are no entities, output nothing."
)

_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
_SPECIAL_TOK_RE = re.compile(r"<\|[^|>]*\|>")


def clean_output(text: str) -> str:
    if not text:
        return ""
    text = _THINK_RE.sub("", text)
    if "</think>" in text:
        text = text.split("</think>")[-1]
    return _SPECIAL_TOK_RE.sub("", text).strip()


def extract_one(sentence: str) -> str:
    """Return the raw model text (entity lines) for one sentence."""
    if not sentence.strip():
        return ""
    kwargs = dict(
        model=QWEN_MODEL,
        messages=[{"role": "system", "content": _SYSTEM},
                  {"role": "user", "content": sentence}],
        temperature=TEMPERATURE,
        max_tokens=MAX_TOKENS,
    )
    if QWEN_NO_THINK:
        kwargs["extra_body"] = {"chat_template_kwargs": {"enable_thinking": False}}
    last_err = None
    for _ in range(MAX_RETRIES):
        try:
            resp = _client.chat.completions.create(**kwargs)
            return clean_output(resp.choices[0].message.content or "")
        except Exception as e:  # transient network / 5xx / rate-limit
            last_err = e
    raise RuntimeError(f"qwen call failed after {MAX_RETRIES} retries: {last_err}")


# ---------------------------------------------------------------------------
# Parse model text -> per-token 21-column tags
# ---------------------------------------------------------------------------
_TYPESET = set(TYPE_ORDER)


def parse_entities(text: str):
    ents = []
    for ln in text.splitlines():
        ln = ln.strip()
        if not ln:
            continue
        parts = re.split(r"\t|\s+-\s+|\s*:\s*|\s*\|\s*", ln)
        if len(parts) >= 2:
            typ = parts[-1].strip().upper()
            surf = " ".join(parts[:-1]).strip().strip("\"'“”«»")
            if typ in _TYPESET and surf:
                ents.append((surf, typ))
    return ents


def entity_tagsets(tokens, ents):
    """Align each entity surface (whitespace tokens) onto `tokens`; per-token set of B-/I- tags."""
    tagsets = [set() for _ in tokens]
    for surf, typ in ents:
        span = surf.split()
        if not span:
            continue
        for i in range(len(tokens) - len(span) + 1):
            if tokens[i:i + len(span)] == span:
                tagsets[i].add(f"B-{typ}")
                for j in range(i + 1, i + len(span)):
                    tagsets[j].add(f"I-{typ}")
    return tagsets


def tags_to_columns(tagset):
    cols = []
    for typ in TYPE_ORDER:
        if f"B-{typ}" in tagset:
            cols.append(f"B-{typ}")
        elif f"I-{typ}" in tagset:
            cols.append(f"I-{typ}")
        else:
            cols.append("O")
    return cols


def read_jsonl(p):
    rows = []
    with io.open(p, encoding="utf-8") as h:
        for line in h:
            if line.strip():
                rows.append(json.loads(line))
    return rows


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    print(f"endpoint : {QWEN_BASE_URL}")
    print(f"model    : {QWEN_MODEL}  (no_think={QWEN_NO_THINK}, temp={TEMPERATURE}, conc={CONCURRENCY})")
    print(f"work dir : {WORK_DIR}")

    src = resolve_konooz_src()
    files = load_domain_files(src)
    data = {os.path.splitext(os.path.basename(f))[0]: read_conll(f) for f in files}
    total = sum(len(v) for v in data.values())
    print(f"data     : {len(files)} domains, {total} sentences ({src})")

    # --- sanity: first 2 sentences ---
    print("\n--- sanity (first 2 sentences) ---")
    dom0 = sorted(data)[0]
    for sent in data[dom0][:2]:
        toks = [r[0] for r in sent]
        raw = extract_one(" ".join(toks))
        print(f"[{dom0}] {' '.join(toks)[:50]!r} -> {parse_entities(raw)[:4]}")
    print()

    # --- resume from cache (keyed by domain, sentence index) ---
    cache = {}
    if os.path.exists(CACHE_FILE):
        for r in read_jsonl(CACHE_FILE):
            cache[(r["domain"], int(r["sidx"]))] = r["text"]
        print(f"[resume] {len(cache)} sentences already cached")

    todo = []
    for dom in sorted(data):
        for sidx, sent in enumerate(data[dom]):
            if (dom, sidx) not in cache:
                todo.append((dom, sidx, " ".join(r[0] for r in sent)))
    print("remaining sentences:", len(todo))

    if todo:
        fh = io.open(CACHE_FILE, "a", encoding="utf-8")
        try:
            with ThreadPoolExecutor(max_workers=CONCURRENCY) as pool:
                futs = {pool.submit(extract_one, s): (d, i) for (d, i, s) in todo}
                for fut in tqdm(as_completed(futs), total=len(futs), desc="extract", unit="sent"):
                    d, i = futs[fut]
                    text = fut.result()
                    cache[(d, i)] = text
                    fh.write(json.dumps({"domain": d, "sidx": i, "text": text}, ensure_ascii=False) + "\n")
                    fh.flush()
        finally:
            fh.close()
    print("total cached:", len(cache))

    # --- build official submission: token + 21 columns, domains concatenated alphabetically ---
    pred_name = f"namaa_pred_{NAME}.txt"
    pred_path = os.path.join(SUBMIT_DIR, pred_name)
    n_ent = 0
    with io.open(pred_path, "w", encoding="utf-8") as f:
        for dom in sorted(data):
            for sidx, sent in enumerate(data[dom]):
                toks = [r[0] for r in sent]
                ents = parse_entities(cache.get((dom, sidx), ""))
                n_ent += len(ents)
                tagsets = entity_tagsets(toks, ents)
                for w, tok in enumerate(toks):
                    f.write(tok + " " + " ".join(tags_to_columns(tagsets[w])) + "\n")
                f.write("\n")
    print(f"predictions -> {pred_path}  (entities aligned: {n_ent})")

    zip_path = os.path.join(SUBMIT_DIR, f"submission_{NAME}.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
        z.write(pred_path, arcname=pred_name)
    print("submission ->", zip_path)
    print("[done] upload submission zip to CodaBench.")


if __name__ == "__main__":
    main()