updated my_model.py

my_model.py

@@ -6,32 +6,22 @@ class StoryPointIncrementModel(nn.Module):
     """
     A custom model wrapper designed to load and use the weights of a fine-tuned 
     Transformer model for regression (story point prediction).
-
-    The missing/unexpected keys error indicates the checkpoint contains the full 
-    Transformer structure. We redefine the model to match that structure.
     """
     
-    def __init__(self, model_name="prajjwal1/bert-tiny", num_labels=1):
+    # CRITICAL FIX: Add cache_dir argument to __init__ and set a default to None
+    def __init__(self, model_name="prajjwal1/bert-tiny", num_labels=1, cache_dir=None):
         super().__init__()
+        
         # Load the configuration of a small BERT-like model as a base template.
-        # The actual weights from model.safetensors will be loaded into this structure.
-        config = AutoConfig.from_pretrained(model_name)
+        # PASS cache_dir to from_pretrained to prevent permission errors
+        config = AutoConfig.from_pretrained(model_name, cache_dir=cache_dir)
         
         # We load the base encoder (up to the pooler)
         self.encoder = AutoModel.from_config(config)
 
-        # The unexpected keys suggest the saved model structure includes a pooler layer.
-        # We define a custom regressor head that will be matched by `load_state_dict` 
-        # (or at least provide a place for the final linear layer if it was saved 
-        # under a different name than the original checkpoint).
-        # We will manually map the final linear layer if necessary.
-        
         # A simple linear layer for regression (predicting a single story point value)
         self.regressor = nn.Linear(config.hidden_size, num_labels)
 
-        # A custom property to track if the loading was successful
-        self.loaded_safetensors_keys = False
-
     def forward(self, input_ids, attention_mask):
         # Pass the tokenized inputs through the Transformer encoder
         outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)