Sravanigunnu/gemma-3-4b-macd-hindi-hate-speech-lora

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

BASE_MODEL = "google/gemma-3-4b-it"
ADAPTER    = "Sravanigunnu/gemma-3-4b-macd-hindi-hate-speech-lora"

tokenizer = AutoTokenizer.from_pretrained(ADAPTER, use_fast=True)
base  = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(base, ADAPTER)
model.eval()

SYSTEM = (
    "You are a hate speech classifier for social media content. "
    "Classify the given text as abusive or non-abusive. "
    "Reply with only '1' if the text contains hate speech or abuse, "
    "or '0' if it is non-abusive. Do not explain your answer."
)

def classify(text: str) -> int:
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user",   "content": text}]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    ) + "Label: "

    ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

    # Resolve token IDs for "0" and "1" in generation context
    base_ids = tokenizer.encode(prompt, add_special_tokens=False)
    id0 = tokenizer.encode(prompt + "0", add_special_tokens=False)[len(base_ids)]
    id1 = tokenizer.encode(prompt + "1", add_special_tokens=False)[len(base_ids)]

    with torch.no_grad():
        logits = model(ids).logits[0, -1, :]

    pair  = torch.stack([logits[id0], logits[id1]])
    probs = torch.softmax(pair, dim=-1)
    return int(probs.argmax().item())   # 0 = non-abusive, 1 = abusive

# Example
print(classify("This is a neutral message."))   # → 0

import pandas as pd
from tqdm import tqdm

def classify_batch(texts: list[str], batch_size: int = 16) -> list[int]:
    tokenizer.padding_side = "left"
    base_ids = tokenizer.encode(
        tokenizer.apply_chat_template(
            [{"role": "system", "content": SYSTEM}, {"role": "user", "content": "x"}],
            tokenize=False, add_generation_prompt=True
        ) + "Label: ",
        add_special_tokens=False,
    )
    sample_prompt = tokenizer.apply_chat_template(
        [{"role": "system", "content": SYSTEM}, {"role": "user", "content": texts[0]}],
        tokenize=False, add_generation_prompt=True
    ) + "Label: "
    id0 = tokenizer.encode(sample_prompt + "0", add_special_tokens=False)[len(base_ids)]
    id1 = tokenizer.encode(sample_prompt + "1", add_special_tokens=False)[len(base_ids)]

    all_preds = []
    for start in tqdm(range(0, len(texts), batch_size)):
        batch = texts[start : start + batch_size]
        prompts = [
            tokenizer.apply_chat_template(
                [{"role": "system", "content": SYSTEM}, {"role": "user", "content": t}],
                tokenize=False, add_generation_prompt=True
            ) + "Label: "
            for t in batch
        ]
        enc = tokenizer(prompts, return_tensors="pt", padding=True,
                        truncation=True, max_length=512).to(model.device)
        with torch.no_grad():
            logits = model(**enc).logits[:, -1, :]
        pair  = torch.stack([logits[:, id0], logits[:, id1]], dim=-1)
        probs = torch.softmax(pair, dim=-1)
        all_preds.extend(probs.argmax(dim=-1).cpu().tolist())
    return all_preds