Update app.py

app.py

@@ -1,43 +1,47 @@
 import os
+import json
 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
+# Remove the datasets import as we won't be using it anymore
+# from datasets import load_dataset
+
 @st.cache_data
 def load_data(file_path):
+    if not os.path.exists(file_path):
+        st.error(f"File not found: {file_path}")
+        return None
     try:
-        return load_dataset('json', data_files=file_path)
+        with open(file_path, 'r') as f:
+            data = json.load(f)
+        return data
     except Exception as e:
         st.error(f"Error loading dataset: {str(e)}")
         return None
 
-@st.cache_resource
-def initialize_model_and_tokenizer(model_name):
-    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, max_length):
-    return tokenizer(examples['prompt'], truncation=True, padding="max_length", max_length=max_length)
+def create_dataset(data, tokenizer, max_length):
+    inputs = []
+    for item in data:
+        prompt = item['prompt']
+        response = item['response']
+        full_text = f"Human: {prompt}\nAssistant: {response}"
+        encoded = tokenizer.encode(full_text, truncation=True, max_length=max_length, padding='max_length')
+        inputs.append(encoded)
+    return inputs
 
 def main():
     st.title("Model Training with Streamlit")
 
-    # User inputs
+    # User inputs with recommended values
     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)
+    file_path = st.text_input("Enter path to training data JSON file", "appointment_training_data.json")
+    max_length = st.number_input("Enter max token length", min_value=32, max_value=512, value=256)
     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")
+    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")
 
     tokenizer, model = initialize_model_and_tokenizer(model_name)
     
@@ -46,18 +50,28 @@ def main():
         return
 
     st.write("Loading and processing dataset...")
-    dataset = load_data(file_path)
+    data = load_data(file_path)
     
-    if dataset is None:
+    if data 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, max_length),
-        batched=True,
-        remove_columns=dataset['train'].column_names
-    )
+    tokenized_dataset = create_dataset(data, tokenizer, max_length)
+
+    # Convert tokenized_dataset to a torch Dataset
+    class SimpleDataset(torch.utils.data.Dataset):
+        def __init__(self, encodings):
+            self.encodings = encodings
+
+        def __getitem__(self, idx):
+            item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
+            return item
+
+        def __len__(self):
+            return len(self.encodings['input_ids'])
+
+    dataset = SimpleDataset({'input_ids': tokenized_dataset})
 
     # Define training arguments
     training_args = TrainingArguments(
@@ -76,7 +90,7 @@ def main():
     trainer = Trainer(
         model=model,
         args=training_args,
-        train_dataset=tokenized_dataset,
+        train_dataset=dataset,
         data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
     )