DanielNRU/Avito-relevance-ruroberta-v3-20260824

import json
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModel, AutoTokenizer

repo_id = "DanielNRU/Avito-relevance-ruroberta-v3-20260824"
config_path = hf_hub_download(repo_id, "config.json")
weights_path = hf_hub_download(repo_id, "model.safetensors")
config = json.load(open(config_path, encoding="utf-8"))

tokenizer = AutoTokenizer.from_pretrained(repo_id)

class AvitoRelevanceClassifier(torch.nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(cfg["model_name"])
        self.dropout1 = torch.nn.Dropout(cfg["dropout_encoder"])
        self.dense = torch.nn.Linear(self.encoder.config.hidden_size, cfg["head_hidden_size"])
        self.activation = torch.nn.GELU()
        self.dropout2 = torch.nn.Dropout(cfg["dropout_head"])
        self.classifier = torch.nn.Linear(cfg["head_hidden_size"], 1)

    def forward(self, input_ids, attention_mask):
        x = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0]
        x = self.dropout1(x)
        x = self.activation(self.dense(x))
        return self.classifier(self.dropout2(x)).squeeze(-1)

model = AvitoRelevanceClassifier(config)
model.load_state_dict(load_file(weights_path, device="cpu"))
model.eval()

text = "Текст для проверки"
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=config["max_length"])
with torch.inference_mode():
    probability = torch.sigmoid(model(**inputs)).item()
relevance = int(probability >= config["decision_threshold_relevant"])
print({"relevance": relevance, "probability_relevant": probability})