Update app.py

app.py

@@ -7,10 +7,9 @@ import torch
 from pytorch_lightning import LightningModule
 from safetensors.torch import save_file
 from torch import nn
-from modelalign import BERTAlignModel 
 
 import gradio as gr
-
+from modelalign import BERTAlignModel
 
 # ===========================
 # Utility Functions
@@ -48,7 +47,11 @@ def load_checkpoint(model: LightningModule, checkpoint_path: str, device: str =
     try:
         # Load the checkpoint; adjust map_location based on device
         checkpoint = torch.load(checkpoint_path, map_location=device)
-        model.load_state_dict(checkpoint['state_dict'], strict=False)
+        # Assuming the checkpoint has a 'state_dict' key
+        if 'state_dict' in checkpoint:
+            model.load_state_dict(checkpoint['state_dict'], strict=False)
+        else:
+            model.load_state_dict(checkpoint, strict=False)
         return True, "Checkpoint loaded successfully."
     except Exception as e:
         return False, f"Failed to load checkpoint: {str(e)}"
@@ -80,31 +83,29 @@ def convert_checkpoint_to_safetensors(checkpoint_url: str, model_name: str):
         # Step 1: Download the checkpoint
         success, message = download_checkpoint(checkpoint_url, checkpoint_path)
         if not success:
-            return gr.update(value=None, visible=False), message
+            return None, message
 
         # Step 2: Initialize the model
         success, model_or_msg = initialize_model(model_name)
         if not success:
-            return gr.update(value=None, visible=False), model_or_msg
+            return None, model_or_msg
         model = model_or_msg
 
         # Step 3: Load the checkpoint
         success, message = load_checkpoint(model, checkpoint_path)
         if not success:
-            return gr.update(value=None, visible=False), message
+            return None, message
 
         # Step 4: Convert to SafeTensors
         success, message = convert_to_safetensors(model, safetensors_path)
         if not success:
-            return gr.update(value=None, visible=False), message
+            return None, message
 
         # Step 5: Read the safetensors file for download
         try:
-            with open(safetensors_path, "rb") as f:
-                safetensors_bytes = f.read()
-            return safetensors_bytes, "Conversion successful! Download your SafeTensors file below."
+            return safetensors_path, "Conversion successful! Download your SafeTensors file below."
         except Exception as e:
-            return gr.update(value=None, visible=False), f"Failed to prepare download: {str(e)}"
+            return None, f"Failed to prepare download: {str(e)}"
 
 # ===========================
 # Gradio Interface Setup
@@ -125,12 +126,20 @@ Convert your PyTorch Lightning `.ckpt` checkpoints to the secure `safetensors` f
 iface = gr.Interface(
     fn=convert_checkpoint_to_safetensors,
     inputs=[
-        gr.inputs.Textbox(lines=2, placeholder="Enter the checkpoint URL here...", label="Checkpoint URL"),
-        gr.inputs.Textbox(lines=1, placeholder="e.g., roberta-base", label="Model Name")
+        gr.Textbox(
+            lines=2, 
+            placeholder="Enter the checkpoint URL here...", 
+            label="Checkpoint URL"
+        ),
+        gr.Textbox(
+            lines=1, 
+            placeholder="e.g., roberta-base", 
+            label="Model Name"
+        )
     ],
     outputs=[
-        gr.outputs.File(label="Download SafeTensors File"),
-        gr.outputs.Textbox(label="Status")
+        gr.File(label="Download SafeTensors File"),
+        gr.Textbox(label="Status")
     ],
     title=title,
     description=description,