Spaces:
Sleeping
Sleeping
| """ | |
| model.py — DistilBERT multi-head model definition and loader | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from transformers import DistilBertModel, DistilBertTokenizer | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class PostWorkoutDistilBERT(nn.Module): | |
| """ | |
| Post-workout multi-head DistilBERT classifier. | |
| 5 independent classification heads sharing one BERT backbone: | |
| - mood (8 classes) | |
| - exertion (3 classes) | |
| - soreness_region (7 classes — which muscle group) | |
| - soreness_severity (4 classes — how intense) | |
| - completion (2 classes) | |
| """ | |
| def __init__( | |
| self, | |
| num_moods: int = 8, | |
| num_exertion_levels: int = 3, | |
| num_soreness_region_classes: int = 7, | |
| num_soreness_severity_classes:int = 4, | |
| num_completion_statuses: int = 2, | |
| ): | |
| super().__init__() | |
| self.bert = DistilBertModel.from_pretrained("distilbert-base-uncased") | |
| hidden_size = self.bert.config.hidden_size # 768 | |
| self.dropout = nn.Dropout(0.3) | |
| self.head_dropout = nn.Dropout(0.1) | |
| # Simple heads for easy tasks | |
| self.mood_head = nn.Linear(hidden_size, num_moods) | |
| self.completion_head = nn.Linear(hidden_size, num_completion_statuses) | |
| # Deeper head for exertion — 768→128→3 | |
| self.exertion_head = nn.Sequential( | |
| nn.Linear(hidden_size, 128), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(128, num_exertion_levels), | |
| ) | |
| # Soreness region head — 768→128→7 (which muscle group) | |
| self.soreness_region_head = nn.Sequential( | |
| nn.Linear(hidden_size, 128), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(128, num_soreness_region_classes), | |
| ) | |
| # Soreness severity head — 768→64→4 (how intense) | |
| self.soreness_severity_head = nn.Sequential( | |
| nn.Linear(hidden_size, 64), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(64, num_soreness_severity_classes), | |
| ) | |
| def forward(self, input_ids, attention_mask): | |
| outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) | |
| cls_output = self.dropout(outputs.last_hidden_state[:, 0, :]) | |
| x = self.head_dropout(cls_output) | |
| return ( | |
| self.mood_head(x), | |
| self.exertion_head(x), | |
| self.soreness_region_head(x), | |
| self.soreness_severity_head(x), | |
| self.completion_head(x), | |
| ) | |
| class PreWorkoutDistilBERT(nn.Module): | |
| """ | |
| Multi-head DistilBERT classifier for pre-workout state analysis. | |
| 6 independent classification heads sharing one BERT backbone: | |
| - mood (8 classes) | |
| - energy (3 classes — low / moderate / high) | |
| - motivation (3 classes — low / moderate / high) | |
| - stress (3 classes — low / moderate / high) | |
| - soreness_region (7 classes — which muscle group) | |
| - soreness_severity (4 classes — how intense) | |
| """ | |
| def __init__( | |
| self, | |
| num_moods: int = 8, | |
| num_energy_levels: int = 3, | |
| num_motivation_levels: int = 3, | |
| num_stress_levels: int = 3, | |
| num_soreness_region_classes: int = 7, | |
| num_soreness_severity_classes:int = 4, | |
| ): | |
| super().__init__() | |
| self.bert = DistilBertModel.from_pretrained("distilbert-base-uncased") | |
| hidden_size = self.bert.config.hidden_size # 768 | |
| self.dropout = nn.Dropout(0.3) | |
| self.head_dropout = nn.Dropout(0.1) | |
| # Simple head — mood has strong linguistic signal | |
| self.mood_head = nn.Linear(hidden_size, num_moods) | |
| # Energy head — 768→128→3 | |
| self.energy_head = nn.Sequential( | |
| nn.Linear(hidden_size, 128), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(128, num_energy_levels), | |
| ) | |
| # Motivation head — 768→64→3 | |
| self.motivation_head = nn.Sequential( | |
| nn.Linear(hidden_size, 64), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(64, num_motivation_levels), | |
| ) | |
| # Stress head — 768→64→3 | |
| self.stress_head = nn.Sequential( | |
| nn.Linear(hidden_size, 64), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(64, num_stress_levels), | |
| ) | |
| # Soreness region head — 768→128→7 | |
| self.soreness_region_head = nn.Sequential( | |
| nn.Linear(hidden_size, 128), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(128, num_soreness_region_classes), | |
| ) | |
| # Soreness severity head — 768→64→4 | |
| self.soreness_severity_head = nn.Sequential( | |
| nn.Linear(hidden_size, 64), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(64, num_soreness_severity_classes), | |
| ) | |
| def forward(self, input_ids, attention_mask): | |
| outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) | |
| cls_output = self.dropout(outputs.last_hidden_state[:, 0, :]) | |
| x = self.head_dropout(cls_output) | |
| return ( | |
| self.mood_head(x), | |
| self.energy_head(x), | |
| self.motivation_head(x), | |
| self.stress_head(x), | |
| self.soreness_region_head(x), | |
| self.soreness_severity_head(x), | |
| ) | |
| def load_pre_model( | |
| model_path: str, | |
| device: torch.device, | |
| num_moods: int = 8, | |
| num_energy_levels: int = 3, | |
| num_motivation_levels: int = 3, | |
| num_stress_levels: int = 3, | |
| num_soreness_region_classes: int = 7, | |
| num_soreness_severity_classes:int = 4, | |
| ): | |
| """ | |
| Instantiate the pre-workout model, load saved weights, set to eval mode. | |
| Returns (model, tokenizer). | |
| """ | |
| logger.info(f"Loading pre-workout model weights from: {model_path}") | |
| model = PreWorkoutDistilBERT( | |
| num_moods=num_moods, | |
| num_energy_levels=num_energy_levels, | |
| num_motivation_levels=num_motivation_levels, | |
| num_stress_levels=num_stress_levels, | |
| num_soreness_region_classes=num_soreness_region_classes, | |
| num_soreness_severity_classes=num_soreness_severity_classes, | |
| ) | |
| state_dict = torch.load(model_path, map_location=device, weights_only=True) | |
| model.load_state_dict(state_dict) | |
| model.to(device) | |
| model.eval() | |
| tokenizer = DistilBertTokenizer.from_pretrained( | |
| "distilbert-base-uncased", | |
| clean_up_tokenization_spaces=True, | |
| ) | |
| logger.info("Pre-workout model loaded and set to eval mode.") | |
| return model, tokenizer | |
| def load_post_model( | |
| model_path: str, | |
| device: torch.device, | |
| num_moods: int = 8, | |
| num_exertion_levels: int = 3, | |
| num_soreness_region_classes: int = 7, | |
| num_soreness_severity_classes:int = 4, | |
| num_completion_statuses: int = 2, | |
| ): | |
| """ | |
| Instantiate the post-workout model, load saved weights, set to eval mode. | |
| Returns (model, tokenizer). | |
| """ | |
| logger.info(f"Loading post-workout model weights from: {model_path}") | |
| model = PostWorkoutDistilBERT( | |
| num_moods=num_moods, | |
| num_exertion_levels=num_exertion_levels, | |
| num_soreness_region_classes=num_soreness_region_classes, | |
| num_soreness_severity_classes=num_soreness_severity_classes, | |
| num_completion_statuses=num_completion_statuses, | |
| ) | |
| state_dict = torch.load(model_path, map_location=device, weights_only=True) | |
| model.load_state_dict(state_dict) | |
| model.to(device) | |
| model.eval() | |
| tokenizer = DistilBertTokenizer.from_pretrained( | |
| "distilbert-base-uncased", | |
| clean_up_tokenization_spaces=True, | |
| ) | |
| logger.info("Post-workout model loaded and set to eval mode.") | |
| return model, tokenizer | |