NOVA-1-prototype / main.py
aledeaux's picture
Create main.py
4d89080 verified
Raw History Blame Contribute Delete
3.89 kB
import tensorflow as tf # type: ignore
from keras_preprocessing.text import Tokenizer # type: ignore
from keras_preprocessing.sequence import pad_sequences # type: ignore
import numpy as np # type: ignore
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
# Optional: suppress ABSEIL warnings if TensorFlow still leaks something
try:
import absl.logging # type: ignore
absl.logging.set_verbosity(absl.logging.ERROR)
except ImportError:
pass
def load_dialogues(datapath="C:/Users/90372/ai/data/", num_words=10000, max_len=40): # noqa
lines_file = os.path.join(datapath, "movie_lines.txt")
convos_file = os.path.join(datapath, "movie_conversations.txt")
id2line = {}
with open(lines_file, encoding="utf-8", errors='ignore') as f:
for line in f:
parts = line.strip().split(" +++$+++ ")
if len(parts) == 5:
line_id, text = parts[0], parts[4]
id2line[line_id] = text
inputs, responses = [], []
with open(convos_file, encoding="utf-8", errors='ignore') as f:
for line in f:
parts = line.strip().split(" +++$+++ ")
if len(parts) == 4:
line_ids = eval(parts[3])
for i in range(len(line_ids) - 1):
input_line = id2line.get(line_ids[i], "")
response_line = id2line.get(line_ids[i + 1], "")
if input_line and response_line:
inputs.append(input_line)
responses.append(response_line)
catchAllError = "I'm sorry, I cant seem to remember anything right now."
tokenizer = Tokenizer(num_words=num_words, oov_token=catchAllError)
tokenizer.fit_on_texts(inputs + responses)
input_sequences = tokenizer.texts_to_sequences(inputs)
response_sequences = tokenizer.texts_to_sequences(responses)
inputs_padded = pad_sequences(input_sequences, maxlen=max_len, padding="post", truncating="post") # noqa
responses_padded = pad_sequences(response_sequences, maxlen=max_len, padding="post", truncating="post") # noqa
return inputs_padded, responses_padded, tokenizer
def build_seq2seq_model(vocab_size, max_len):
inputs = tf.keras.Input(shape=(max_len,))
embedding = tf.keras.layers.Embedding(vocab_size, 256)(inputs)
lstm = tf.keras.layers.LSTM(256, return_sequences=True)(embedding)
attention = tf.keras.layers.Attention()([lstm, lstm])
decoder_lstm = tf.keras.layers.LSTM(256)(attention)
outputs = tf.keras.layers.Dense(vocab_size)(decoder_lstm) # Final output layer
model = tf.keras.Model(inputs, outputs)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) # Removed label_smoothing
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4), loss=loss) # Compile model
return model
def train_conversational_model():
inputs, responses, tokenizer = load_dialogues()
target_sequences = responses[:, 1:] # Shift the target by one timestep
target_sequences = np.pad(target_sequences, ((0, 0), (0, 1)), mode='constant') # Pad the target sequences
model = build_seq2seq_model(vocab_size=10000, max_len=inputs.shape[1])
model.fit(inputs, target_sequences, epochs=5, batch_size=32) # Train with shifted target sequences
return model, tokenizer
def generate_response(model, tokenizer, prompt, max_len=40):
sequence = tokenizer.texts_to_sequences([prompt])
sequence = pad_sequences(sequence, maxlen=max_len)
pred = model.predict(sequence)
next_word_id = np.argmax(pred[0])
return tokenizer.index_word.get(next_word_id, "<UNK>")
if __name__ == "__main__":
model, tokenizer = train_conversational_model()
while True:
prompt = input("User: ")
if prompt.lower() == "exit":
break
print("AI:", generate_response(model, tokenizer, prompt))