| """Search queries for passages, Qwen3.8-27B via vLLM (Apache-2.0 outputs). |
| IN=passages.parquet SHARD=i NSHARD=n OUT=dir python rerank/gen_queries.py -> $OUT/q_<shard>.jsonl {pid, query, qtype, qlang} |
| Per passage: one natural-language question and one short keyword-style query, in the passage's language; for 15% of |
| non-English passages one more query in English (cross-lingual search). Queries must be answerable by the passage but |
| should not copy long phrases from it. |
| """ |
| import json, os, random, re, sys |
| sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "sentiment")) |
|
|
|
|
| def parse_json(t): |
| t = re.sub(r"^```(?:json)?|```$", "", t.strip(), flags=re.M).strip() |
| s, e = t.find("{"), t.rfind("}") |
| try: |
| return json.loads(t[s: e + 1]) |
| except Exception: |
| return None |
|
|
|
|
| def prompt(text, xl): |
| extra = ('\n- "english": a natural question in English that this passage answers (for cross-lingual search)' if xl else "") |
| return f"""Passage: |
| <<< |
| {text} |
| >>> |
| |
| Write search queries that a user might type into a search engine and for which this passage is a good result. |
| - "question": a natural question in the same language as the passage, answered by the passage |
| - "keywords": a short keyword-style query (2-6 words) in the same language as the passage{extra} |
| Do not copy long phrases from the passage; use the words a searcher would use. Return only JSON with these keys.""" |
|
|
|
|
| if __name__ == "__main__": |
| import pandas as pd |
| from vllm import LLM, SamplingParams |
| from common import hard_exit |
| df = pd.read_parquet(os.environ["IN"]); sh, ns = int(os.environ.get("SHARD", 0)), int(os.environ.get("NSHARD", 1)) |
| df = df.iloc[sh::ns] |
| rng = random.Random(11 + sh) |
| xl = [l != "eng_Latn" and rng.random() < 0.15 for l in df.lang] |
| llm = LLM(os.environ.get("MODEL", "Qwen/Qwen3.8-27B-FP8"), max_model_len=4096, gpu_memory_utilization=0.92, max_num_seqs=512, |
| limit_mm_per_prompt={"image": 0, "video": 0}) |
| outs = llm.chat([[{"role": "user", "content": prompt(t, x)}] for t, x in zip(df.text, xl)], |
| SamplingParams(temperature=0.7, top_p=0.95, max_tokens=300, seed=5 + sh), chat_template_kwargs={"enable_thinking": False}) |
| os.makedirs(os.environ["OUT"], exist_ok=True); n = 0 |
| with open(f"{os.environ['OUT']}/q_{sh}.jsonl", "w") as f: |
| for pid, lang, o in zip(df.pid, df.lang, outs): |
| j = parse_json(o.outputs[0].text) |
| if not isinstance(j, dict): |
| continue |
| for qt in ["question", "keywords", "english"]: |
| q = j.get(qt) |
| if isinstance(q, str) and 3 <= len(q.strip()) <= 300: |
| f.write(json.dumps(dict(pid=int(pid), query=q.strip(), qtype=qt, qlang="eng_Latn" if qt == "english" else lang), ensure_ascii=False) + "\n"); n += 1 |
| print("queries", n, "for", len(df), "passages", flush=True) |
| hard_exit(0) |
|
|
