LLM for political classification

classify_sample.py

import argparse
import os
import re

cache_root = os.path.abspath(
    "YOUR_PATH_HERE"
)
os.environ["HF_HOME"] = cache_root
os.environ["HF_HUB_CACHE"] = os.path.join(cache_root, "hub")
os.environ["VLLM_CACHE_ROOT"] = os.path.join(cache_root, "vllm")

import pandas as pd
import torch
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from vllm import LLM, SamplingParams


PROMPT_PATH = "zero_4_v4.prompt"


def make_tsv_safe(text):
    return str(text).replace("\t", " ").replace("\n", " ").replace("\r", " ")


def load_prompt(prompt_path):
    with open(prompt_path, "r", encoding="utf-8") as f:
        return f.read()


def load_tokenizer_with_fallback(model_tag: str):
    tokenizer_kwargs = {}
    if model_tag == "nvidia/Llama-3_1-Nemotron-51B-Instruct":
        tokenizer_kwargs["trust_remote_code"] = True
    try:
        return AutoTokenizer.from_pretrained(model_tag, **tokenizer_kwargs)
    except AttributeError as exc:
        if "'list' object has no attribute 'keys'" in str(exc):
            print(
                "Fast tokenizer initialization failed due to special token config. "
                "Retrying with use_fast=False."
            )
            return AutoTokenizer.from_pretrained(
                model_tag, use_fast=False, **tokenizer_kwargs
            )
        raise


def parse_answer(text):
    text = str(text).strip()
    match = re.search(r"\b([123])\b", text)
    if match:
        return match.group(1)
    return None


def ensure_nemotron_transformers_compat():
    from transformers.generation import utils as generation_utils

    if not hasattr(generation_utils, "NEED_SETUP_CACHE_CLASSES_MAPPING"):
        generation_utils.NEED_SETUP_CACHE_CLASSES_MAPPING = {}


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("model_tag")
    parser.add_argument(
        "--data_path",
        default="YOUR_DATA_PATH",
    )
    parser.add_argument("--out_prefix", default="dolma")
    parser.add_argument("--prompt_path", default=PROMPT_PATH)
    return parser.parse_args()


def main():
    args = parse_args()

    model_tag = args.model_tag
    data_path = args.data_path
    out_prefix = args.out_prefix
    prompt_template = load_prompt(args.prompt_path)
    tokenizer = load_tokenizer_with_fallback(model_tag)
    sampling_params = SamplingParams(temperature=0.0, max_tokens=16)
    num_gpus = torch.cuda.device_count()
    max_model_len = 8192
    use_transformers_backend = False
    hf_model = None

    def build_prompt(text):
        return prompt_template.format(content=text)

    def load_nemotron_transformers_8bit():
        ensure_nemotron_transformers_compat()
        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
        device_map = "auto" if num_gpus > 1 else "cuda:0"
        return AutoModelForCausalLM.from_pretrained(
            model_tag,
            trust_remote_code=True,
            quantization_config=bnb_config,
            device_map=device_map,
            torch_dtype=torch.bfloat16,
        )

    def get_hf_input_device(model):
        if hasattr(model, "hf_device_map") and model.hf_device_map:
            for device in model.hf_device_map.values():
                if device in ("cpu", "disk"):
                    continue
                if isinstance(device, int):
                    return torch.device(f"cuda:{device}")
                return torch.device(str(device))
        return next(model.parameters()).device

    if model_tag == "nvidia/Llama-3_1-Nemotron-51B-Instruct":
        llm_kwargs = {
            "model": model_tag,
            "max_model_len": max_model_len,
            "quantization": "bitsandbytes",
            "load_format": "bitsandbytes",
            "trust_remote_code": True,
        }
        if num_gpus > 1:
            llm_kwargs["pipeline_parallel_size"] = num_gpus
        try:
            llm = LLM(**llm_kwargs)
        except Exception as exc:
            print(
                "vLLM initialization failed for Nemotron "
                f"({exc.__class__.__name__}: {exc}). "
                "Falling back to Transformers 8-bit bitsandbytes backend."
            )
            use_transformers_backend = True
            hf_model = load_nemotron_transformers_8bit()
    elif "bnb" in model_tag:
        llm = LLM(
            model=model_tag,
            pipeline_parallel_size=num_gpus,
            max_model_len=max_model_len,
        )
    else:
        llm = LLM(
            model=model_tag, tensor_parallel_size=num_gpus, max_model_len=max_model_len
        )

    def model_call(user_text):
        user_prompt = build_prompt(user_text)
        messages = [{"role": "user", "content": user_prompt}]
        chat_template_kwargs = {"tokenize": False, "add_generation_prompt": True}
        if "Qwen3.5" in model_tag:
            chat_template_kwargs["enable_thinking"] = False
        chat_input = tokenizer.apply_chat_template(messages, **chat_template_kwargs)
        if use_transformers_backend:
            model_inputs = tokenizer(chat_input, return_tensors="pt")
            input_device = get_hf_input_device(hf_model)
            model_inputs = {k: v.to(input_device) for k, v in model_inputs.items()}
            with torch.inference_mode():
                generated = hf_model.generate(
                    **model_inputs,
                    do_sample=False,
                    max_new_tokens=16,
                    pad_token_id=tokenizer.eos_token_id,
                )
            prompt_len = model_inputs["input_ids"].shape[-1]
            return tokenizer.decode(
                generated[0][prompt_len:], skip_special_tokens=True
            ).strip()
        outputs = llm.generate([chat_input], sampling_params, use_tqdm=False)
        return outputs[0].outputs[0].text.strip()

    if data_path.endswith(".txt"):
        with open(data_path, "r", encoding="utf-8") as f:
            texts = [line.strip() for line in f if line.strip()]
        df = pd.DataFrame({"text": texts})
    else:
        df = pd.read_json(data_path, lines=True, compression="gzip")
    model_name_safe = model_tag.replace("/", "_")
    out_path = f"{out_prefix}_classification_result_{model_name_safe}_zero_4_v4.txt"

    def count_existing_output_lines(path: str) -> int:
        if not os.path.exists(path):
            return 0
        count = 0
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            for _ in f:
                count += 1
        return count

    already_done = count_existing_output_lines(out_path)
    if already_done > 0:
        print(f"Resuming: found {already_done} existing lines in {out_path}")
    if already_done >= len(df):
        if already_done > len(df):
            print(
                f"Warning: output has {already_done} lines but input has {len(df)} rows. "
                "Nothing to do."
            )
        return

    with open(out_path, "a", encoding="utf-8") as out:
        for idx, text in enumerate(tqdm(df.text, total=len(df))):
            if idx < already_done:
                continue
            if pd.isnull(text) or not str(text).strip():
                out.write("Empty input")
                out.write("\t")
                out.write("\n")
                continue
            try:
                annotation = model_call(str(text))
                answer = parse_answer(annotation)
                if answer:
                    out.write(make_tsv_safe(answer))
                else:
                    out.write("No answer")
            except Exception:
                out.write("Error")
            finally:
                out.write("\t")
                out.write(make_tsv_safe(text))
                out.write("\n")


if __name__ == "__main__":
    main()