Update README.md

README.md

@@ -61,11 +61,36 @@ classification of news categories politics, society and conflicts.
 Example of how to use the model:
 
 ```python
+import torch.nn as nn
+from transformers import BertModel
 import torch
 from transformers import AutoTokenizer
 from huggingface_hub import hf_hub_download
 
 
+class BiLSTMClassifier(nn.Module):
+    def __init__(self, hidden_dim, output_dim, n_layers, dropout):
+        super(BiLSTMClassifier, self).__init__()
+        self.bert = BertModel.from_pretrained("bert-base-multilingual-cased")
+        self.lstm = nn.LSTM(self.bert.config.hidden_size, hidden_dim, num_layers=n_layers, 
+                            bidirectional=True, dropout=dropout, batch_first=True)
+        self.fc = nn.Linear(hidden_dim * 2, output_dim)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, input_ids, attention_mask, labels=None):
+            with torch.no_grad():
+                embedded = self.bert(input_ids=input_ids, attention_mask=attention_mask)[0]
+            lstm_out, _ = self.lstm(embedded)
+            pooled = torch.mean(lstm_out, dim=1)
+            logits = self.fc(self.dropout(pooled))
+            
+            if labels is not None:
+                loss_fn = nn.CrossEntropyLoss()
+                loss = loss_fn(logits, labels)
+                return {"loss": loss, "logits": logits}  # Возвращаем словарь
+            return logits  # Возвращаем логиты, если метки не переданы
+
+
 categories = ['climate', 'conflicts', 'culture', 'economy', 'gloss', 'health',
               'politics', 'science', 'society', 'sports', 'travel']