Update app.py

app.py

@@ -2,13 +2,10 @@ import os
 import json
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
-from transformers import TextDataset, DataCollatorForLanguageModeling
+from transformers import DataCollatorForLanguageModeling
 import torch
 from tqdm import tqdm
 
-# 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):
@@ -22,6 +19,16 @@ def load_data(file_path):
         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 create_dataset(data, tokenizer, max_length):
     inputs = []
     for item in data:
@@ -65,13 +72,12 @@ def main():
             self.encodings = encodings
 
         def __getitem__(self, idx):
-            item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
-            return item
+            return torch.tensor(self.encodings[idx])
 
         def __len__(self):
-            return len(self.encodings['input_ids'])
+            return len(self.encodings)
 
-    dataset = SimpleDataset({'input_ids': tokenized_dataset})
+    dataset = SimpleDataset(tokenized_dataset)
 
     # Define training arguments
     training_args = TrainingArguments(