Update app.py

app.py

@@ -2,60 +2,100 @@ import os
 import streamlit as st
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+from transformers import TextDataset, DataCollatorForLanguageModeling
+import torch
+from tqdm import tqdm
 
 # Streamlit caching functions
 @st.cache_data
-def load_data():
-    # Load the dataset
-    return load_dataset('json', data_files='training_data.json')
+def load_data(file_path):
+    try:
+        return load_dataset('json', data_files=file_path)
+    except Exception as e:
+        st.error(f"Error loading dataset: {str(e)}")
+        return None
 
 @st.cache_resource
 def initialize_model_and_tokenizer(model_name):
-    # Initialize the tokenizer and model
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    return tokenizer, model
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        return tokenizer, model
+    except Exception as e:
+        st.error(f"Error initializing model and tokenizer: {str(e)}")
+        return None, None
 
-def preprocess_function(examples, tokenizer):
-    return tokenizer(examples['prompt'], truncation=True, padding="max_length", max_length=128)
+def preprocess_function(examples, tokenizer, max_length):
+    return tokenizer(examples['prompt'], truncation=True, padding="max_length", max_length=max_length)
 
 def main():
     st.title("Model Training with Streamlit")
-    
-    model_name = 'mistral/Mixtral-8x7B'  # Replace with the correct model name
+
+    # User inputs
+    model_name = st.text_input("Enter model name", "distilgpt2")
+    file_path = st.text_input("Enter path to training data JSON file", "training_data.json")
+    max_length = st.number_input("Enter max token length", min_value=32, max_value=512, value=128)
+    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=4)
+    learning_rate = st.number_input("Enter learning rate", min_value=1e-6, max_value=1e-3, value=2e-5, format="%.1e")
+
     tokenizer, model = initialize_model_and_tokenizer(model_name)
+    
+    if tokenizer is None or model is None:
+        st.warning("Failed to initialize model and tokenizer. Please check the model name and try again.")
+        return
 
     st.write("Loading and processing dataset...")
-    dataset = load_data()
+    dataset = load_data(file_path)
     
-    # Preprocess the dataset
+    if dataset is None:
+        st.warning("Failed to load dataset. Please check the file path and try again.")
+        return
+
     st.write("Tokenizing dataset...")
-    tokenized_dataset = dataset['train'].map(lambda x: preprocess_function(x, tokenizer), batched=True)
+    tokenized_dataset = dataset['train'].map(
+        lambda x: preprocess_function(x, tokenizer, max_length),
+        batched=True,
+        remove_columns=dataset['train'].column_names
+    )
 
     # Define training arguments
     training_args = TrainingArguments(
-        output_dir='./results',              # Output directory
-        evaluation_strategy='epoch',         # Evaluation strategy
-        learning_rate=2e-5,                  # Learning rate
-        per_device_train_batch_size=4,       # Batch size for training
-        per_device_eval_batch_size=4,        # Batch size for evaluation
-        num_train_epochs=3,                  # Number of training epochs
-        weight_decay=0.01,                   # Strength of weight decay
-        logging_dir='./logs',                # Directory for storing logs
-        logging_steps=10,                    # Log every 10 steps
+        output_dir='./results',
+        evaluation_strategy='epoch',
+        learning_rate=learning_rate,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=batch_size,
+        num_train_epochs=num_epochs,
+        weight_decay=0.01,
+        logging_dir='./logs',
+        logging_steps=10,
     )
-    
+
     # Initialize the Trainer
     trainer = Trainer(
-        model=model,                         # The model to train
-        args=training_args,                  # Training arguments
-        train_dataset=tokenized_dataset,     # Training dataset
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
     )
-    
+
     if st.button('Start Training'):
         st.write("Starting training...")
-        trainer.train()
+        progress_bar = st.progress(0)
+        
+        for epoch in range(int(num_epochs)):
+            trainer.train()
+            progress = (epoch + 1) / num_epochs
+            progress_bar.progress(progress)
+            
+            # Save the model after each epoch
+            model_path = f"./results/model_epoch_{epoch+1}"
+            trainer.save_model(model_path)
+            st.write(f"Model saved: {model_path}")
+
         st.write("Training complete.")
+        st.write("You can now use the trained model for inference or further fine-tuning.")
 
 if __name__ == "__main__":
-    main()
+    main()