Upload 2 files

ebert_model.py

@@ -0,0 +1,40 @@
+from transformers import BertConfig, BertModel
+import torch.nn as nn
+
+class EBertConfig(BertConfig):
+    model_type = "ebert"
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.adapter_size = kwargs.pop('adapter_size', None)
+
+class EBertModel(BertModel):
+    config_class = EBertConfig
+
+    def __init__(self, config: EBertConfig):
+        super().__init__(config)
+        if config.adapter_size:
+            self.adapters = nn.ModuleList([
+                nn.Sequential(
+                    nn.Linear(config.hidden_size, config.adapter_size),
+                    nn.ReLU(),
+                    nn.Linear(config.adapter_size, config.hidden_size),
+                )
+                for _ in range(config.num_hidden_layers)
+            ])
+        else:
+            self.adapters = None
+
+    def forward(self, *args, **kwargs):
+        outputs = super().forward(*args, **kwargs)
+        sequence_output = outputs.last_hidden_state
+
+        if self.adapters is not None:
+            for adapter in self.adapters:
+                sequence_output = sequence_output + adapter(sequence_output)
+
+        return outputs.__class__(
+            last_hidden_state=sequence_output,
+            pooler_output=outputs.pooler_output,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

proverka.py

@@ -0,0 +1,66 @@
+from transformers import BertTokenizerFast, BertConfig
+import torch
+from datasets import load_dataset
+from ebert_model import EBertConfig, EBertModel
+from transformers import BertForMaskedLM
+from safetensors.torch import load_file
+import os
+
+model_path = "./ebert"
+tokenizer = BertTokenizerFast.from_pretrained(model_path)
+
+config = EBertConfig.from_pretrained(model_path)
+model = BertForMaskedLM(config)
+model.bert = EBertModel(config)
+
+weights_path = f"{model_path}/model.safetensors"
+if os.path.exists(weights_path):
+    state_dict = load_file(weights_path)
+    model.load_state_dict(state_dict, strict=False)
+else:
+    raise FileNotFoundError(f"Файл {weights_path} не найден.")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
+dataset = load_dataset("Expotion/russian-facts-qa", split="train")
+
+def predict_masked_text(example):
+    text = f"{example['q'].strip()} {example['a'].strip()}"
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, padding="max_length")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    input_ids = inputs["input_ids"].clone()
+    labels = input_ids.clone()
+    mask_token_index = 2
+    if input_ids.size(1) <= mask_token_index:
+        return {
+            "original_text": text,
+            "masked_text": "Слишком короткий текст",
+            "predicted_tokens": [],
+            "true_token": ""
+        }
+    input_ids[0, mask_token_index] = tokenizer.mask_token_id
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predictions = outputs.logits
+    predicted_token_ids = torch.topk(predictions[0, mask_token_index], 5).indices.tolist()
+    predicted_tokens = [tokenizer.decode([id]).strip() for id in predicted_token_ids]
+
+    return {
+        "original_text": text,
+        "masked_text": tokenizer.decode(input_ids[0], skip_special_tokens=True),
+        "predicted_tokens": predicted_tokens,
+        "true_token": tokenizer.decode([labels[0, mask_token_index]], skip_special_tokens=True)
+    }
+
+total_params = sum(p.numel() for p in model.parameters())
+print(f"Общее количество параметров: {total_params}")
+
+num_examples = 1
+for i in range(num_examples):
+    result = predict_masked_text(dataset[i])
+    print(f"Оригинал: {result['original_text']}")
+    print(f"Замаскированный: {result['masked_text']}")
+    print(f"Предсказание: {result['predicted_tokens']}")
+    print(f"Истина: {result['true_token']}")