Updated model.py to work with AutoModel and trust_remote_code

model.py

@@ -96,6 +96,7 @@ class Block(nn.Module):
 
 class GPT(PreTrainedModel):
     config_class = GPTConfig
+    base_model_prefix = "gpt"
 
     def __init__(self, config):
         super().__init__(config)
@@ -135,13 +136,13 @@ class GPT(PreTrainedModel):
             x = block(x)
         x = self.transformer.ln_f(x)
 
+        logits = self.lm_head(x)
+
         if labels is not None:
-            logits = self.lm_head(x)
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-1)
-            return {'logits': logits, 'loss': loss}
+            return {"loss": loss, "logits": logits}
         else:
-            logits = self.lm_head(x[:, [-1], :])
-            return {'logits': logits}
+            return {"logits": logits}
 
     @torch.no_grad()
     def generate(self, input_ids, max_new_tokens, temperature=1.0, top_k=None):