Spaces:
Running on Zero
Running on Zero
File size: 2,380 Bytes
957c899 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """
Custom TextCNN model made Hugging Face Hub compatible.
Architecture inferred from checkpoint weights:
- Embedding: (30522 vocab, 140 dim)
- 3 parallel Conv1d branches: kernel sizes 3, 4, 5 — each with 100 output channels
- Final FC layer: (300 -> 1) [300 = 100*3 concatenated conv outputs]
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin
class TextCNNConfig:
def __init__(
self,
vocab_size: int = 30522,
embed_dim: int = 140,
num_filters: int = 100,
kernel_sizes: list = [3, 4, 5],
num_classes: int = 1,
**kwargs,
):
self.vocab_size = vocab_size
self.embed_dim = embed_dim
self.num_filters = num_filters
self.kernel_sizes = kernel_sizes
self.num_classes = num_classes
class TextCNN(nn.Module, PyTorchModelHubMixin):
"""
TextCNN for MCQ scoring, pushable to / loadable from the Hugging Face Hub.
Usage:
# Save + push
model = TextCNN(config)
model.load_state_dict(torch.load("original_state_dict.pt"))
model.push_to_hub("your-username/cnn-mcq-model")
# Load back
model = TextCNN.from_pretrained("your-username/cnn-mcq-model")
"""
def __init__(self, config: dict = None):
super().__init__()
if config is None:
config = {}
cfg = TextCNNConfig(**config)
self.config = config if config else cfg.__dict__
self.embedding = nn.Embedding(cfg.vocab_size, cfg.embed_dim)
self.convs = nn.ModuleList([
nn.Conv1d(in_channels=cfg.embed_dim, out_channels=cfg.num_filters, kernel_size=k)
for k in cfg.kernel_sizes
])
self.fc = nn.Linear(cfg.num_filters * len(cfg.kernel_sizes), cfg.num_classes)
def forward(self, input_ids):
# input_ids: (batch, seq_len)
x = self.embedding(input_ids) # (batch, seq_len, embed_dim)
x = x.permute(0, 2, 1) # (batch, embed_dim, seq_len)
conv_outs = [F.relu(conv(x)) for conv in self.convs]
pooled = [F.max_pool1d(c, c.size(2)).squeeze(2) for c in conv_outs]
cat = torch.cat(pooled, dim=1) # (batch, num_filters * len(kernel_sizes))
out = self.fc(cat) # (batch, num_classes)
return out
|