Views
No views yet
1import json, torch
2from transformers import AutoTokenizer
3
4with open("config.json") as f:
5 cfg = json.load(f)
6
7all_tags = cfg["all_tags"]
8tokenizer = AutoTokenizer.from_pretrained("DanielNRU/S7-tags-ruroberta-20260811")
9# model = RuRoBERTaTagClassifier("ai-forever/ruRoBERTa-large", num_tags=len(all_tags))
10# model.load_state_dict(torch.load("best_pytorch_model.bin", map_location="cpu"))
11model.eval()
12
13text = "Потеряли багаж, служба поддержки не отвечает"
14enc = tokenizer(
15 text,
16 max_length=512,
17 truncation=True,
18 padding="max_length",
19 return_tensors="pt",
20)
21with torch.no_grad():
22 logits = model(enc["input_ids"], enc["attention_mask"])
23 probs = torch.sigmoid(logits).squeeze(0).tolist()
24
25result = {
26 all_tags[i]: round(probs[i], 3)
27 for i in range(len(all_tags))
28 if probs[i] >= 0.3
29}
30print(result)