File size: 1,809 Bytes
50da3a7 |
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 |
# Save this as model.py
import torch
import torch.nn as nn
from sentence_transformers import SentenceTransformer
class SummaryClassifier(nn.Module):
def __init__(self, embedder, num_classes, dropout=0.1):
"""
Initializes the classifier.
Args:
embedder: A pre-loaded SentenceTransformer model.
num_classes (int): The number of output classes.
dropout (float): Dropout probability.
"""
super().__init__()
self.embedder = embedder
embedding_dim = embedder.get_sentence_embedding_dimension()
self.head = nn.Sequential(
nn.Dropout(dropout),
nn.Linear(embedding_dim, 128),
nn.ReLU(),
nn.Linear(128, num_classes)
)
# Freeze the embedder parameters
for p in self.embedder.parameters():
p.requires_grad = False
def forward(self, texts, return_embeddings=False):
"""
Forward pass.
Args:
texts (list[str]): A list of input strings.
return_embeddings (bool): Whether to return embeddings alongside logits.
Returns:
torch.Tensor: The output logits.
(Optional) torch.Tensor: The sentence embeddings.
"""
# Automatically use the same device as the model's 'head'
target_device = next(self.head.parameters()).device
embeddings = self.embedder.encode(
texts,
convert_to_tensor=True,
show_progress_bar=False,
device=str(target_device)
)
logits = self.head(embeddings)
if return_embeddings:
return logits, embeddings
return logits |