Update finetune.py

finetune.py

@@ -172,9 +172,12 @@ model.state_dict = (
 import gradio as gr
 
 def train():
+    print(os.listdir(OUTPUT_DIR))
     # Call your trainer's train() function here
     trainer.train()
     print("Training complete.")  # optional message to display when training is done
+    model.save_pretrained(OUTPUT_DIR)
+    repo.push_to_hub(OUTPUT_DIR, commit_message="Ft model")
 
 iface = gr.Interface(
     fn=train,
@@ -188,7 +191,9 @@ iface = gr.Interface(
     allow_screenshot=False
 )
 
-model.save_pretrained(OUTPUT_DIR)
-repo.push_to_hub(OUTPUT_DIR, commit_message="Ft model")
+if __name__ == '__main__':
+    iface.launch()
+
+