Update app.py

app.py

@@ -13,10 +13,10 @@ def train_model(data_file, model_name, epochs, batch_size, learning_rate):
         
         df = pd.read_csv(data_file.name)
         
-        if 'prompt' not in df.columns or 'description' not in df.columns:
-            return "CSV file must contain 'prompt' and 'description' columns.", None
+        if 'prompt' not in df.columns or 'completion' not in df.columns:
+            return "CSV file must contain 'prompt' and 'completion' columns.", None
         
-        df['text'] = df['prompt'] + ': ' + df['description']
+        df['text'] = df['prompt'] + ': ' + df['completion']
         dataset = Dataset.from_pandas(df[['text']])
         
         tokenizer = GPT2Tokenizer.from_pretrained(model_name)