CategorAI / train.py
noch inator
added the AI files
8e8f2a5
Raw
History Blame Contribute Delete
2.89 kB
import pickle
import tensorflow as tf
import pandas as pd
import re
# Define the function to preprocess the input text
def preprocess_text(text):
# Convert the text to lowercase
text = text.lower()
# Replace any non-alphanumeric characters with a space
text = re.sub(r'[^a-zA-Z0-9\'+\-*/\s]', ' ', text)
# Replace the word 'plus' with the '+' symbol
text = text.replace('plus', '+')
# Replace the word 'minus' with the '-' symbol
text = text.replace('minus', '-')
# Replace the word 'times' with the '*' symbol
text = text.replace('times', '*')
# Replace the word 'divided by' with the '/' symbol
text = text.replace('divided by', '/')
# Replace the characters * / - + with a space and the character itself
text = text.replace('*', ' * ').replace('/', ' / ').replace('-', ' - ').replace('+', ' + ').replace('.', ' . ')
# Replace all numbers with a special token
text = re.sub(r'\d+', '<num>', text)
# Replace all numbers with a special token
text = text.replace(r'oov', '<oov>')
# Remove any extra whitespace
text = re.sub(r'\s+', ' ', text)
return text.strip()
epochs = int(input("epochs: "))
batch_size = int(input("batch size: "))
# Load the training data from a CSV file
data = pd.read_csv("data.csv")
input_data = data['input'].values
labels = data['class'].values
# Preprocess the input data
input_data = [preprocess_text(text) for text in input_data]
# Define the vocabulary and encode the input data
vocab = tf.keras.preprocessing.text.Tokenizer(filters='')
vocab.fit_on_texts(input_data)
encoded_input = vocab.texts_to_sequences(input_data)
# Pad the encoded input data to ensure all inputs are of the same length
max_length = max([len(seq) for seq in encoded_input])
padded_input = tf.keras.preprocessing.sequence.pad_sequences(encoded_input, maxlen=max_length, padding='post')
# Define the neural network model
model = tf.keras.Sequential([
tf.keras.layers.Embedding(input_dim=len(vocab.word_index) + 1, output_dim=64, input_length=max_length),
tf.keras.layers.Conv1D(filters=64, kernel_size=3, activation='relu', padding='same'),
tf.keras.layers.GlobalMaxPooling1D(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(128, activation=tf.keras.layers.ELU(alpha=1.0)),
tf.keras.layers.Dense(3, activation='softmax')
])
# Compile the model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
# Train the model
model.fit(padded_input, labels, epochs=epochs, batch_size=batch_size)
# Save the trained model
model.save('main.h5')
# Save the tokenizer using pickle
with open('tokenizer.pkl', 'wb') as file:
pickle.dump(vocab, file)
token_to_word = {token: word for word, token in vocab.word_index.items()}
print(token_to_word)
model.summary()
# Enter into an interactive loop to test the model