Spaces:
Running
Running
File size: 939 Bytes
5b60104 86f6d4d 5b60104 86f6d4d 5b60104 86f6d4d | 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 | import torch
import torch.nn as nn
class ProgrammingClassifier(nn.Module):
"""
Classifies programming-related intent or domain.
"""
def __init__(self, input_dim: int = 128, num_classes: int = 5):
super().__init__()
# store attributes for engine access
self.input_dim = input_dim
self.num_classes = num_classes
# classifier layer
self.classifier = nn.Linear(input_dim, num_classes)
# label names
self.labels = [
"general",
"debugging",
"architecture",
"algorithm",
"deployment"
]
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.classifier(x)
def predict_label(self, logits: torch.Tensor) -> str:
"""
Convert logits to programming domain label
"""
idx = torch.argmax(logits, dim=-1).item()
return self.labels[idx] |