Update app.py

app.py

@@ -4,7 +4,11 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
 from transformers import DataCollatorForLanguageModeling
 import torch
-from tqdm import tqdm
+from huggingface_hub import Repository, HfFolder
+
+# Authenticate Hugging Face Hub
+hf_token = st.secrets["hf_token"]  # Store your token in the Hugging Face Space Secrets
+HfFolder.save_token(hf_token)
 
 @st.cache_data
 def load_data(file_path):
@@ -64,6 +68,9 @@ def main():
     num_epochs = st.number_input("Enter number of training epochs", min_value=1, max_value=10, value=3)
     batch_size = st.number_input("Enter batch size", min_value=1, max_value=32, value=8)
     learning_rate = st.number_input("Enter learning rate", min_value=1e-6, max_value=1e-3, value=5e-5, format="%.1e")
+    
+    # Specify your Hugging Face model repository ID
+    repo_id = st.text_input("Enter Hugging Face repository ID", "nileshhanotia/PeVe")
 
     tokenizer, model = initialize_model_and_tokenizer(model_name)
     
@@ -118,19 +125,25 @@ def main():
     if st.button('Start Training'):
         st.write("Starting training...")
         progress_bar = st.progress(0)
+
+        # Initialize the Hugging Face repository
+        repo = Repository(local_dir="./results", clone_from=repo_id)
         
         for epoch in range(int(num_epochs)):
             trainer.train()
             progress = (epoch + 1) / num_epochs
             progress_bar.progress(progress)
             
-            # Save the model after each epoch
+            # Save the model after each epoch locally
             model_path = f"./results/model_epoch_{epoch+1}"
             trainer.save_model(model_path)
-            st.write(f"Model saved: {model_path}")
+            st.write(f"Model saved locally: {model_path}")
+
+            # Push to Hugging Face Hub
+            repo.push_to_hub(commit_message=f"Model after epoch {epoch+1}")
+            st.write(f"Model pushed to Hugging Face Hub: {repo_id}")
 
-        st.write("Training complete.")
-        st.write("You can now use the trained model for inference or further fine-tuning.")
+        st.write("Training complete. Model is available on the Hugging Face Hub.")
 
 if __name__ == "__main__":
-    main()
+    main()