Training in progress, step 75000, checkpoint

config.json

@@ -2,7 +2,7 @@
   "attention_bias": false,
   "attention_dropout": 0.0,
   "auto_map": {
-    "AutoModel": "modeling_novomolgen.NovoMolGen"
+    "AutoModelForCausalLM": "modeling_novomolgen.NovoMolGen"
   },
   "bos_token_id": 2,
   "eos_token_id": 3,

modeling_novomolgen.py

@@ -98,7 +98,7 @@ class NovoMolGenConfig(LlamaConfig):
         self.fused_dropout_add_ln = fused_dropout_add_ln
         self.residual_in_fp32 = residual_in_fp32
         self.loss_type = loss_type
-        self.auto_map = {"AutoModel": "modeling_novomolgen.NovoMolGen"}
+        self.auto_map = {"AutoModelForCausalLM": "modeling_novomolgen.NovoMolGen"}
 
     @classmethod
     def from_pretrained(
@@ -328,3 +328,7 @@ class NovoMolGen(GPTLMHeadModel):
                 eos_position = eos_position[0, 0].item()  # Get the index of the first occurrence
                 output[i, eos_position + 1:] = eos_id
         return output
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        # HF’s GenerationMixin would normally do more, but for a basic LM this usually suffices:
+        return {"input_ids": input_ids, "attention_mask": attention_mask}