Views
No views yet
1import torch.nn as nn
2from transformers import BertModel
3import torch
4from transformers import AutoTokenizer
5from huggingface_hub import hf_hub_download
6
7
8class BiLSTMClassifier(nn.Module):
9 def __init__(self, hidden_dim, output_dim, n_layers, dropout):
10 super(BiLSTMClassifier, self).__init__()
11 self.bert = BertModel.from_pretrained("bert-base-multilingual-cased")
12 self.lstm = nn.LSTM(self.bert.config.hidden_size, hidden_dim, num_layers=n_layers,
13 bidirectional=True, dropout=dropout, batch_first=True)
14 self.fc = nn.Linear(hidden_dim * 2, output_dim)
15 self.dropout = nn.Dropout(dropout)
16
17 def forward(self, input_ids, attention_mask, labels=None):
18 with torch.no_grad():
19 embedded = self.bert(input_ids=input_ids, attention_mask=attention_mask)[0]
20 lstm_out, _ = self.lstm(embedded)
21 pooled = torch.mean(lstm_out, dim=1)
22 logits = self.fc(self.dropout(pooled))
23
24 if labels is not None:
25 loss_fn = nn.CrossEntropyLoss()
26 loss = loss_fn(logits, labels)
27 return {"loss": loss, "logits": logits} # Возвращаем словарь
28 return logits # Возвращаем логиты, если метки не переданы
29
30
31categories = ['climate', 'conflicts', 'culture', 'economy', 'gloss', 'health',
32 'politics', 'science', 'society', 'sports', 'travel']
33
34repo_id = "data-silence/lstm-news-classifier"
35tokenizer = AutoTokenizer.from_pretrained(repo_id)
36model_path = hf_hub_download(repo_id=repo_id, filename="model.pth")
37
38model = torch.load(model_path)
39
40def get_predictions(news: str, model) -> str:
41 with torch.no_grad():
42 inputs = tokenizer(news, return_tensors="pt")
43 del inputs['token_type_ids']
44 output = model.forward(**inputs)
45 id_best_label = torch.argmax(output[0, :], dim=-1).detach().cpu().numpy()
46 prediction = categories[id_best_label]
47 return prediction
48
49
50# Использование классификатора
51get_predictions('В Париже завершилась церемония завершения Олимпийский игр', model=model)
52# 'sports'