1from transformers import AutoTokenizer, AutoModelForCausalLM
2from peft import PeftModel
3import torch
4
5BASE_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
6ADAPTER = "Sravanigunnu/llama-3.1-8b-macd-hindi-n2000-lora"
7
8tokenizer = AutoTokenizer.from_pretrained(ADAPTER, use_fast=True)
9base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16, device_map="auto")
10model = PeftModel.from_pretrained(base, ADAPTER)
11model.eval()
12
13SYSTEM = (
14 "You are a hate speech classifier for social media content. "
15 "Classify the given text as abusive or non-abusive. "
16 "Reply with only '1' if the text contains hate speech or abuse, "
17 "or '0' if it is non-abusive. Do not explain your answer."
18)
19
20def classify(text: str) -> int:
21 messages = [{"role": "system", "content": SYSTEM},
22 {"role": "user", "content": text}]
23 prompt = tokenizer.apply_chat_template(
24 messages, tokenize=False, add_generation_prompt=True
25 ) + "Label: "
26
27 base_ids = tokenizer.encode(prompt, add_special_tokens=False)
28 id0 = tokenizer.encode(prompt + "0", add_special_tokens=False)[len(base_ids)]
29 id1 = tokenizer.encode(prompt + "1", add_special_tokens=False)[len(base_ids)]
30
31 ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
32 with torch.no_grad():
33 logits = model(ids).logits[0, -1, :]
34
35 probs = torch.softmax(torch.stack([logits[id0], logits[id1]]), dim=-1)
36 return int(probs.argmax().item()) # 0 = non-abusive, 1 = abusive