binhquoc/alm-add-rq1-ent

import warnings
warnings.simplefilter("ignore")
import torch.nn.functional as F
import numpy as np
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import os
import json
from dotenv import load_dotenv
load_dotenv()

os.environ["WANDB_DISABLED"] = "true"

num_labels = 6

model_name = "Qwen/Qwen3-8B"
lora_path = "binhquoc/alm-add-rq1-ent"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(model_name,
                                          trust_remote_code=True,
                                          cache_dir=os.environ["CACHE_DIR"])

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

base_model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=num_labels,
    device_map="auto",
    problem_type="multi_label_classification",
    trust_remote_code=True,
    quantization_config=bnb_config,
    cache_dir=os.environ["CACHE_DIR"]
)

base_model.config.pad_token_id = tokenizer.pad_token_id

model = PeftModel.from_pretrained(base_model, lora_path)
model.eval()