Upload app.py

app.py

@@ -40,13 +40,13 @@ def load_model():
     vocab_size = 50257
     
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = SmolLM(vocab_size, embed_dim, num_heads, num_layers, max_seq_len).to(device)
+    #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)
-    #config = checkpoint['config']
-    #model = GPT(config)
-    #model.load_state_dict(checkpoint['model_state_dict'])
+    #model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
+    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_state_dict'])
     model.to(device)
     model.eval()  # Set to evaluation mode