ArdLi/ToxiTrace

import torch
from torch import nn
import torch.nn.functional as F
from transformers import BertTokenizer

class EndClassifier(nn.Module):
    def __init__(self, dropout, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dense1 = nn.Linear(768, 2)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)


    def forward(self, x):
        return self.dense1(x)

class BertClissifier(nn.Module):
    def __init__(self, bert, tokenizer, dropout=0.1, freeze_fine_tune=True, freeze_last_layer=True, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.tokenizer = tokenizer
        self.bert = bert
        self.freeze = freeze_fine_tune
        if freeze_fine_tune:
            for name, param in self.bert.named_parameters():
                if "encoder.layer.11" in name and not freeze_last_layer:
                    param.requires_grad_(self.freeze)
                else:
                    param.requires_grad_(False)
        self.ffn = EndClassifier(dropout)

    def forward(self, input_ids=None, attention_mask=None, **kwargs):
        bert_out = self.bert(input_ids, attention_mask=attention_mask).pooler_output
        logits = self.ffn(bert_out)
        return F.softmax(logits, dim=1)

    def logits_from_ids(self, input_ids, attention_mask=None):
        pool = self.bert(input_ids=input_ids, attention_mask=attention_mask).pooler_output
        return self.ffn(pool)

    def logits_from_embeds(self, inputs_embeds, attention_mask=None):
        out = self.bert(inputs_embeds=inputs_embeds, attention_mask=attention_mask)
        pool = out.pooler_output
        return self.ffn(pool)

    def set_freeze(self, freeze: bool):
        self.freeze = freeze
        for _, param in self.bert.named_parameters():
            param.requires_grad_(not self.freeze)

    @property
    def is_freeze(self):
        return self.freeze

device = "cuda" if torch.cuda.is_available() else "cpu"
model = torch.load("ToxiTrace_COLD_RoBERTa.pt", map_location=device, weights_only=False)
model.eval()

tokenizer = BertTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext")

text = "我觉得你们河南人都是坏人"
enc = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
input_ids = enc["input_ids"].to(device)
attention_mask = enc["attention_mask"].to(device)

with torch.no_grad():
    logits = model.logits_from_ids(input_ids, attention_mask)
    probs = F.softmax(logits, dim=-1)
    pred = probs.argmax(dim=-1).item()

label_map = {0: "Non-offensive", 1: "Offensive"}
print(f"Prediction : {label_map[pred]}")
print(f"Confidence : {probs.squeeze(0).tolist()}")