Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import json
 from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments, AutoTokenizer
 from datasets import Dataset
-import torch
+import shutil
+import os
 
 # Load tokenized data
 def load_data(tokenized_file):
@@ -51,7 +52,10 @@ def fine_tune_model(tokenized_file, progress=gr.Progress()):
     model.save_pretrained('./fine_tuned_patentbert')
     tokenizer.save_pretrained('./fine_tuned_patentbert')
     
-    return "Model fine-tuned and saved successfully."
+    # Create a zip file of the fine-tuned model
+    shutil.make_archive('fine_tuned_patentbert', 'zip', './fine_tuned_patentbert')
+    
+    return "Model fine-tuned and saved successfully. Download the model using the link below.", "fine_tuned_patentbert.zip"
 
 # Create Gradio interface
 iface = gr.Interface(
@@ -59,7 +63,10 @@ iface = gr.Interface(
     inputs=[
         gr.File(label="Upload Tokenized Data JSON")
     ],
-    outputs=gr.Textbox(label="Processing Information"),
+    outputs=[
+        gr.Textbox(label="Processing Information"),
+        gr.File(label="Download Fine-Tuned Model")
+    ],
     title="Fine-Tune Patent BERT Model",
     description="Upload tokenized JSON file to fine-tune the BERT model.",
     live=True  # Enable live updates for progress