Upload app.py

app.py

@@ -32,7 +32,7 @@ class SmolLM(nn.Module):
 
 
 def load_model():
-    checkpoint_path = 'smollm_checkpoint.pth'
+    checkpoint_path = 'final_checkpoint.pth'
     embed_dim = 512
     num_heads = 8
     num_layers = 4
@@ -43,11 +43,12 @@ def load_model():
     #model = SmolLM(vocab_size, embed_dim, num_heads, num_layers, max_seq_len).to(device)
     
     #model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
-    checkpoint = torch.load(checkpoint_path, map_location=device)
+    #checkpoint = torch.load(checkpoint_path, map_location=device)
     #config = checkpoint['config']
-    model = SmolLM(vocab_size, embed_dim, num_heads, num_layers, max_seq_len)
-    model.load_state_dict(checkpoint)
-    model.to(device)
+    #model = SmolLM(vocab_size, embed_dim, num_heads, num_layers, max_seq_len)
+    #model.load_state_dict(checkpoint)
+    model = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    model = model.to(device)
     model.eval()  # Set to evaluation mode
     
     # Disable gradient computation