/
maratgaliulin
/
landcode_classifier
Обзор
Документация
Войти
/
maratgaliulin
/
landcode_classifier
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
CI/CD
Аналитика
Безопасность
master
methods/classes/BertWithNumericAndCluster.py
35 строк
1 KB
maratgaliulin
.
23 июн 2026, 15:04
23 июн 2026, 15:04
488c495
Код
Авторство
О чём код?
# methods/classes/BertWithNumericAndCluster.py import torch CLUSTER_EMBED_DIM = 32 class BertWithNumericAndCluster(torch.nn.Module): """ BERT + area cluster embedding (+ optional scaled numeric area) → single classifier head. Used for legacy flat checkpoints; new training uses BertWithNumericAndClusterHierarchical. """ def __init__(self, bert_model, num_labels, num_clusters=6, dropout=0.1, use_numeric_area=False): super().__init__() self.bert = bert_model self.dropout = torch.nn.Dropout(dropout) self.use_numeric_area = use_numeric_area self.num_clusters = num_clusters self.cluster_embedding = torch.nn.Embedding(num_clusters, CLUSTER_EMBED_DIM) extra = CLUSTER_EMBED_DIM + (1 if use_numeric_area else 0) self.classifier = torch.nn.Linear(self.bert.config.hidden_size + extra, num_labels) def _shared_repr(self, input_ids, attention_mask, cluster_ids, area=None): outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) pooled_output = self.dropout(outputs.pooler_output) cluster_embeds = self.cluster_embedding(cluster_ids) parts = [pooled_output, cluster_embeds] if self.use_numeric_area: parts.append(area.unsqueeze(1)) return torch.cat(parts, dim=1) def forward(self, input_ids, attention_mask, cluster_ids, area=None): combined = self._shared_repr(input_ids, attention_mask, cluster_ids, area) return self.classifier(combined)