ritaberrada/iol-bnb-test

1import os2os.environ["HF_HUB_OFFLINE"] = "1"3os.environ["TRANSFORMERS_OFFLINE"] = "1"4MODEL_ID = "."5 6import json7import pandas as pd8import torch9from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig10 11# bitsandbytes 4-bit (NF4). float16 compute dtype: the T4 is Turing, no native bfloat16.12bnb_config = BitsAndBytesConfig(13    load_in_4bit=True,14    bnb_4bit_quant_type="nf4",15    bnb_4bit_compute_dtype=torch.float16,16)17 18tok = AutoTokenizer.from_pretrained(MODEL_ID)19model = AutoModelForCausalLM.from_pretrained(20    MODEL_ID, quantization_config=bnb_config, device_map="auto"21).eval()22print("loaded model in 4-bit (bitsandbytes)", flush=True)23 24df = pd.read_csv("/tmp/data/test.csv", dtype=str).fillna("")25 26rows = []27for _, r in df.iterrows():28    messages = [29        {"role": "system", "content":30            "You solve International Linguistics Olympiad problems. Answer every numbered "31            "item. Put each answer on its own line, in order, with no numbering and no extra text."},32        {"role": "user", "content": f"{r['context'].strip()}\n\n{r['query'].strip()}"},33    ]34    ids = tok.apply_chat_template(35        messages, add_generation_prompt=True, return_tensors="pt",36    ).to(model.device)37    with torch.no_grad():38        out = model.generate(ids, max_new_tokens=256, do_sample=False)39    text = tok.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()40    answers = [ln.strip() for ln in text.splitlines() if ln.strip()]41    rows.append({"id": r["id"], "pred": json.dumps(answers, ensure_ascii=False)})42    print(f"{len(rows)}/{len(df)} done", flush=True)43 44pd.DataFrame(rows).to_csv("submission.csv", index=False)45print("wrote submission.csv", flush=True)46