LLM for political classification

classify_sample_upd.py

import argparse
import csv
import gzip
import json
import os
import re
import sys

cache_root = os.path.abspath(
    "/mnt/hum01-home01/y79782dn/rds/dominik-llama/extremism_detection/hf_cache"
)
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 torch
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from vllm import LLM, SamplingParams


PROMPT_PATH = "zero_4_v4.prompt"


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="/mnt/hum01-rds/Nikolaev_Dmitry/corpora/dolma/sample.jsonl.gz",
    )
    parser.add_argument("--out_prefix", default="dolma")
    parser.add_argument("--prompt_path", default=PROMPT_PATH)
    parser.add_argument(
        "--label_field",
        default="label",
        help="Name of the field added to each output record with the classification label.",
    )
    parser.add_argument(
        "--part",
        choices=("all", "first", "second"),
        default="all",
        help="Process all documents (default), or only the first or second half.",
    )
    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()

    is_gzip_csv = data_path.endswith(".csv.gz")
    csv_fieldnames = ["id", "text", "qwen35", "generated_response", "url"]

    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()]
        records = [{"text": text} for text in texts]
    elif is_gzip_csv:
        try:
            csv.field_size_limit(sys.maxsize)
        except OverflowError:
            csv.field_size_limit(2**31 - 1)
        with gzip.open(data_path, "rt", encoding="utf-8", newline="") as f:
            records = [
                {field: row.get(field, "") for field in csv_fieldnames}
                for row in csv.DictReader(f)
            ]
    else:
        open_input = gzip.open if data_path.endswith(".gz") else open
        with open_input(data_path, "rt", encoding="utf-8") as f:
            records = [json.loads(line) for line in f if line.strip()]

    if args.part == "first":
        records = records[: len(records) // 2]
    elif args.part == "second":
        records = records[len(records) // 2 :]

    model_name_safe = model_tag.replace("/", "_")
    part_suffix = "" if args.part == "all" else f"_{args.part}"
    out_path_base = (
        f"{out_prefix}_classification_result_{model_name_safe}_zero_4_v4"
        f"{part_suffix}"
    )
    out_path = f"{out_path_base}.csv.gz" if is_gzip_csv else f"{out_path_base}.jsonl"

    def count_existing_output_lines(path: str) -> int:
        if not os.path.exists(path):
            return 0
        if is_gzip_csv:
            with gzip.open(path, "rt", encoding="utf-8", newline="") as f:
                return sum(1 for _ in csv.DictReader(f))
        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(records):
        if already_done > len(records):
            print(
                f"Warning: output has {already_done} lines but input has {len(records)} rows. "
                "Nothing to do."
            )
        return

    if is_gzip_csv:
        output_fieldnames = csv_fieldnames + [args.label_field]
        open_output = gzip.open(out_path, "at", encoding="utf-8", newline="")
    else:
        open_output = open(out_path, "a", encoding="utf-8")

    with open_output as out:
        csv_writer = None
        if is_gzip_csv:
            csv_writer = csv.DictWriter(out, fieldnames=output_fieldnames)
            if already_done == 0:
                csv_writer.writeheader()
        for idx, record in enumerate(tqdm(records, total=len(records))):
            if idx < already_done:
                continue
            output_record = dict(record)
            text = record.get("text")
            if text is None or not str(text).strip():
                label = "Empty input"
            else:
                try:
                    annotation = model_call(str(text))
                    label = parse_answer(annotation) or "No answer"
                except Exception:
                    label = "Error"
            output_record[args.label_field] = label
            if csv_writer:
                csv_writer.writerow(output_record)
            else:
                out.write(json.dumps(output_record, ensure_ascii=False, default=str))
                out.write("\n")


if __name__ == "__main__":
    main()