Create train.py

train.py

@@ -0,0 +1,57 @@
+import argparse
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.layers import LSTM, Dense
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.callbacks import ModelCheckpoint
+from datasets import load_dataset
+
+# Parse command-line arguments
+parser = argparse.ArgumentParser(description='Train a text generation model.')
+parser.add_argument('filename', help='Name of the text file to train on')
+parser.add_argument('-s', '--save', help='Name of the file to save the trained model')
+parser.add_argument('-e', '--epochs', type=int, default=100, help='Number of training epochs')
+args = parser.parse_args()
+
+dataset = load_dataset("code_search_net", "python")
+# Load the text file
+with open(args.filename, 'r', encoding='utf-8') as f:
+    text = f.read()
+
+# Create character mappings
+chars = sorted(list(set(text)))
+char_to_index = {char: index for index, char in enumerate(chars)}
+index_to_char = {index: char for index, char in enumerate(chars)}
+
+# Prepare the training data
+max_sequence_length = 100
+step = 1
+sequences = []
+next_chars = []
+for i in range(0, len(text) - max_sequence_length, step):
+    sequences.append(text[i:i+max_sequence_length])
+    next_chars.append(text[i+max_sequence_length])
+X = np.zeros((len(sequences), max_sequence_length, len(chars)), dtype=np.bool)
+y = np.zeros((len(sequences), len(chars)), dtype=np.bool)
+for i, sequence in enumerate(sequences):
+    for t, char in enumerate(sequence):
+        X[i, t, char_to_index[char]] = 1
+    y[i, char_to_index[next_chars[i]]] = 1
+
+# Define the model
+model = Sequential([
+    LSTM(128, input_shape=(max_sequence_length, len(chars))),
+    Dense(len(chars), activation='softmax')
+])
+model.compile(loss='categorical_crossentropy', optimizer='adam')
+
+# Train the model
+if args.save:
+    checkpoint = ModelCheckpoint(args.save, monitor='loss', verbose=1, save_best_only=True, mode='min')
+    callbacks_list = [checkpoint]
+    history = model.fit(X, y, batch_size=128, epochs=args.epochs, callbacks=callbacks_list)
+else:
+    history = model.fit(X, y, batch_size=128, epochs=args.epochs)
+
+# Print the final loss
+print(f"Final loss: {history.history['loss'][-1]:.4f}")