Upload app.py

app.py

@@ -184,9 +184,36 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
 model = SmolLM2(SmolLM2Config())
 
-# Load trained weights
-checkpoint = torch.load('checkpoint_step_5000.pt', map_location=device)
-model.load_state_dict(checkpoint['model_state_dict'])
+# Load trained weights from zip
+def load_checkpoint_from_zip(zip_path, device):
+    try:
+        with zipfile.ZipFile(zip_path, 'r') as zf:
+            with zf.open('checkpoint.pt') as f:
+                buffer = io.BytesIO(f.read())
+                checkpoint = torch.load(buffer, map_location=device)
+                if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                    return checkpoint['model_state_dict']
+                return checkpoint
+    except Exception as e:
+        print(f"Error loading checkpoint: {e}")
+        return None
+
+# Try to load the checkpoint
+try:
+    checkpoint_path = 'checkpoint_step_5000.zip'  # or .pt file
+    if checkpoint_path.endswith('.zip'):
+        state_dict = load_checkpoint_from_zip(checkpoint_path, device)
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        state_dict = checkpoint['model_state_dict'] if isinstance(checkpoint, dict) else checkpoint
+    
+    if state_dict is not None:
+        model.load_state_dict(state_dict)
+        print("Checkpoint loaded successfully")
+except Exception as e:
+    print(f"Error loading checkpoint: {e}")
+    print("Initializing model with random weights")
+
 model.to(device)
 model.eval()