Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,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