Upload app.py

app.py

@@ -2,6 +2,8 @@ import torch
 from transformers import AutoTokenizer
 from model import SmolLM2, SmolLM2Config
 import gradio as gr
+import zipfile
+import io
 
 # Initialize model and tokenizer
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -9,7 +11,8 @@ tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
 model = SmolLM2(SmolLM2Config())
 
 # Load trained weights
-checkpoint = torch.load('checkpoints/checkpoint_step_5000.pt', map_location=device)  # Adjust path as needed
+
+checkpoint = torch.load('checkpoint_step_5000.pt', map_location=device)  # Adjust path as needed
 model.load_state_dict(checkpoint['model_state_dict'])
 model.to(device)
 model.eval()