Update app.py

app.py

@@ -1,62 +1,73 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+import os
 from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
+import pandas as pd
+from io import StringIO
 
-# Streamlit App title
-st.title("Fine-Tune Mixtral 8x7B Model")
-
-# Model name input field
-model_name = "mistral/mixtral-8x7b"
-
-# Access the Hugging Face token from Streamlit secrets
-token = st.secrets["HF_TOKEN"]
-
-# Load the tokenizer and model
-try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
-    model = AutoModelForCausalLM.from_pretrained(model_name, token=token)
-    st.write("Model and tokenizer loaded successfully!")
-except Exception as e:
-    st.error(f"An error occurred while loading the model: {e}")
-
-# Load the dataset from the existing file
-dataset_path = "training_data.json"
-
-try:
-    dataset = load_dataset('json', data_files={'train': dataset_path})
-except Exception as e:
-    st.error(f"An error occurred while loading the dataset: {e}")
-
-# Tokenize the dataset
 def preprocess_function(examples):
+    if 'prompt' not in examples:
+        raise ValueError("Key 'prompt' not found in examples. Please check your dataset fields.")
     return tokenizer(examples['prompt'], truncation=True, padding="max_length", max_length=128)
 
-tokenized_dataset = dataset['train'].map(preprocess_function, batched=True)
+def train_model(training_data):
+    # Load the dataset
+    dataset = load_dataset('json', data_files={'train': training_data})
+    
+    # Initialize the tokenizer and model
+    model_name = 'mistral/Mixtral-8x7B'  # Replace with the correct model name
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Training arguments for fine-tuning
-training_args = TrainingArguments(
-    output_dir="./results",
-    eval_strategy="epoch",  # Use `eval_strategy` instead of `evaluation_strategy`
-    learning_rate=2e-5,
-    per_device_train_batch_size=1,
-    num_train_epochs=3,
-    weight_decay=0.01,
-)
+    # Tokenize the dataset
+    tokenized_dataset = dataset['train'].map(preprocess_function, batched=True)
+    
+    # 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
+    )
+    
+    # Initialize the Trainer
+    trainer = Trainer(
+        model=model,                         # The model to train
+        args=training_args,                  # Training arguments
+        train_dataset=tokenized_dataset,     # Training dataset
+    )
+    
+    # Start training
+    trainer.train()
 
-# Initialize the Trainer for fine-tuning
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_dataset,
-)
+def main():
+    st.title("Model Training with Streamlit")
+    
+    st.write("Upload your training data in JSON format:")
+    uploaded_file = st.file_uploader("Choose a file", type="json")
+    
+    if uploaded_file is not None:
+        st.write("File uploaded successfully!")
+        
+        # Read the file into a pandas DataFrame
+        file_contents = uploaded_file.read().decode("utf-8")
+        st.write("Preview of uploaded data:")
+        st.text(file_contents[:1000])  # Display first 1000 characters for preview
+        
+        # Save the file to a temporary location
+        temp_file_path = 'training_data.json'
+        with open(temp_file_path, 'w') as f:
+            f.write(file_contents)
+        
+        # Call the train_model function
+        st.write("Training the model...")
+        train_model(temp_file_path)
+        st.write("Training completed!")
 
-# Button to start fine-tuning
-if st.button("Start Fine-Tuning"):
-    with st.spinner("Fine-tuning in progress..."):
-        try:
-            trainer.train()
-            st.success("Fine-tuning completed!")
-            model.save_pretrained("./fine_tuned_model")
-            st.write("Fine-tuned model saved!")
-        except Exception as e:
-            st.error(f"An error occurred during fine-tuning: {e}")
+if __name__ == "__main__":
+    main()