Update app.py

app.py

@@ -1,15 +1,155 @@
+import os
+import requests
+import tempfile
+import shutil
+
+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
 
-def get_safetensors():
-    with open("AlignScore-base.safetensors", "rb") as f:
-        return f.read()
+
+# ===========================
+# Utility Functions
+# ===========================
+
+def download_checkpoint(url: str, dest_path: str):
+    """
+    Downloads the checkpoint from the specified URL to the destination path.
+    """
+    try:
+        with requests.get(url, stream=True) as response:
+            response.raise_for_status()
+            with open(dest_path, 'wb') as f:
+                shutil.copyfileobj(response.raw, f)
+        return True, "Checkpoint downloaded successfully."
+    except Exception as e:
+        return False, f"Failed to download checkpoint: {str(e)}"
+
+def initialize_model(model_name: str, device: str = 'cpu'):
+    """
+    Initializes the BERTAlignModel based on the provided model name.
+    """
+    try:
+        model = BERTAlignModel(base_model_name=model_name)
+        model.to(device)
+        model.eval()  # Set to evaluation mode
+        return True, model
+    except Exception as e:
+        return False, f"Failed to initialize model: {str(e)}"
+
+def load_checkpoint(model: LightningModule, checkpoint_path: str, device: str = 'cpu'):
+    """
+    Loads the checkpoint into the model.
+    """
+    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)
+        return True, "Checkpoint loaded successfully."
+    except Exception as e:
+        return False, f"Failed to load checkpoint: {str(e)}"
+
+def convert_to_safetensors(model: LightningModule, save_path: str):
+    """
+    Converts the model's state_dict to the safetensors format.
+    """
+    try:
+        state_dict = model.state_dict()
+        save_file(state_dict, save_path)
+        return True, "Model converted to SafeTensors successfully."
+    except Exception as e:
+        return False, f"Failed to convert to SafeTensors: {str(e)}"
+
+# ===========================
+# Gradio Interface Function
+# ===========================
+
+def convert_checkpoint_to_safetensors(checkpoint_url: str, model_name: str):
+    """
+    Orchestrates the download, loading, conversion, and preparation for download.
+    Returns the safetensors file or an error message.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        checkpoint_path = os.path.join(tmpdir, "model.ckpt")
+        safetensors_path = os.path.join(tmpdir, "model.safetensors")
+
+        # Step 1: Download the checkpoint
+        success, message = download_checkpoint(checkpoint_url, checkpoint_path)
+        if not success:
+            return gr.update(value=None, visible=False), 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
+        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
+
+        # Step 4: Convert to SafeTensors
+        success, message = convert_to_safetensors(model, safetensors_path)
+        if not success:
+            return gr.update(value=None, visible=False), 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."
+        except Exception as e:
+            return gr.update(value=None, visible=False), f"Failed to prepare download: {str(e)}"
+
+# ===========================
+# Gradio Interface Setup
+# ===========================
+
+title = "Checkpoint to SafeTensors Converter"
+description = """
+Convert your PyTorch Lightning `.ckpt` checkpoints to the secure `safetensors` format.
+
+**Inputs**:
+- **Checkpoint URL**: Direct link to the `.ckpt` file.
+- **Model Name**: Name of the base model (e.g., `roberta-base`, `bert-base-uncased`).
+
+**Output**:
+- Downloadable `safetensors` file.
+"""
 
 iface = gr.Interface(
-    fn=get_safetensors,
-    inputs=[],
-    outputs=gr.outputs.File(label="Download SafeTensors Model"),
-    title="Download SafeTensors Model",
-    description="Click the button below to download the SafeTensors version of the model."
+    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")
+    ],
+    outputs=[
+        gr.outputs.File(label="Download SafeTensors File"),
+        gr.outputs.Textbox(label="Status")
+    ],
+    title=title,
+    description=description,
+    examples=[
+        [
+            "https://huggingface.co/yzha/AlignScore/resolve/main/AlignScore-base.ckpt?download=true",
+            "roberta-base"
+        ],
+        [
+            "https://path.to/your/checkpoint.ckpt",
+            "bert-base-uncased"
+        ]
+    ],
+    allow_flagging="never"
 )
 
-iface.launch()
+# ===========================
+# Launch the Interface
+# ===========================
+
+if __name__ == "__main__":
+    iface.launch()