Update app.py

app.py

@@ -237,11 +237,18 @@ tokenizer = CharTokenizer()
 # =========================================================
 # 2. Model Definition (CUSTOM – REQUIRED)
 # =========================================================
-class TransformerCopyConfig(PretrainedConfig):
-    model_type = "transformer_copy"
-    def __init__(self, vocab_size=tokenizer.vocab_size, **kwargs):
-        super().__init__(**kwargs)
-        self.vocab_size = vocab_size
+class TransformerCopyHF(PreTrainedModel):
+    config_class = TransformerCopyConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = TransformerCopyModel(
+            vocab_size=config.vocab_size,
+            d_model=256,
+            nhead=8,
+            num_layers=4,
+            dim_ff=512
+        )
+
 
 class PositionalEncoding(nn.Module):
     def __init__(self, d_model, max_len=512):
@@ -264,15 +271,25 @@ class TransformerCopyModel(nn.Module):
         self.embedding = nn.Embedding(vocab_size, d_model)
         self.pos = PositionalEncoding(d_model)
 
-        enc_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_ff, dropout=dropout)
-        dec_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_ff, dropout=dropout)
-
-        self.encoder = nn.TransformerEncoder(enc_layer, num_layers)
-        self.decoder = nn.TransformerDecoder(dec_layer, num_layers)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_ff,
+            dropout=dropout
+        )
+        decoder_layer = nn.TransformerDecoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_ff,
+            dropout=dropout
+        )
 
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
+        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers)
         self.fc = nn.Linear(d_model, vocab_size)
 
 
+
     def forward(self, src, tgt):
         src_emb = self.pos(self.embedding(src))
         tgt_emb = self.pos(self.embedding(tgt))