import streamlit as st import tensorflow as tf from keras.layers import Input, Dense, Embedding, MultiHeadAttention from keras.layers import Dropout, LayerNormalization from keras.models import Model from keras.utils import pad_sequences from tensorflow.keras import layers import numpy as np import logging logging.getLogger('tensorflow').setLevel(logging.ERROR) class TransformerChatbot(Model): def __init__(self, vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate): super(TransformerChatbot, self).__init__() self.embedding = Embedding(vocab_size, d_model) self.attention = MultiHeadAttention(num_heads=n_head, key_dim=d_model) self.norm1 = LayerNormalization(epsilon=1e-6) self.dropout1 = Dropout(dropout_rate) self.dense1 = Dense(ff_dim, activation="relu") self.dense2 = Dense(d_model) self.norm2 = LayerNormalization(epsilon=1e-6) self.dropout2 = Dropout(dropout_rate) self.flatten = tf.keras.layers.Flatten() self.fc = Dense(vocab_size, activation="softmax") self.max_len = max_len def call(self, inputs): x = self.embedding(inputs) # Masking mask = self.create_padding_mask(inputs) attn_output = self.attention(x, x, x, attention_mask=mask) x = x + attn_output x = self.norm1(x) x = self.dropout1(x) x = self.dense1(x) x = self.dense2(x) x = self.norm2(x) x = self.dropout2(x) x = self.fc(x) return x def create_padding_mask(self, seq): mask = tf.cast(tf.math.equal(seq, 0), tf.float32) return mask[:, tf.newaxis, tf.newaxis, :] def completion_model(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate,weights,datafile,dict,len2,text2): with open(datafile,"r") as f: text = f.read() text = text.lower() words = text.split() loaded_dict = np.load(dict, allow_pickle=True) word_to_num = loaded_dict["word_to_num"].item() num_to_word = loaded_dict["num_to_word"].item() X = [] Y = [] for i in range(len(words)-1): word = words[i] next_word = words[i+1] X.append(word_to_num[word]) Y.append(word_to_num[next_word]) Y.append(0) X.append(word_to_num[words[-1]]) X_train = pad_sequences([X]) y_train = pad_sequences([Y]) chatbot = TransformerChatbot(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate) chatbot.load_weights(weights) chatbot.build(input_shape=(None, max_len)) # Build the model chatbot.compile(optimizer="adam", loss="sparse_categorical_crossentropy") for i in range(1): other_text2 = text2 other_text2 = other_text2.lower() other_words2 = other_text2.split() other_num2 = [word_to_num[word] for word in other_words2] given_X2 = other_num2 input_sequence2 = pad_sequences([given_X2], maxlen=max_len, padding='post') output_sentence = other_text2 + "" for _ in range(len2): predicted_token = np.argmax(chatbot.predict(input_sequence2), axis=-1) predicted_token = predicted_token.item() out = num_to_word[predicted_token] # if out == ".": # break output_sentence += " " + out given_X2 = given_X2[1:] given_X2.append(predicted_token) input_sequence2 = pad_sequences([given_X2], maxlen=max_len, padding='post') out2 = output_sentence return out2 st.title("UniGLM TEXT completion Model") st.subheader("Next Word Prediction AI Model by Webraft-AI") #Picking what NLP task you want to do option = st.selectbox('Model',('13M_OLD','26M_OLD')) #option is stored in this variable #Textbox for text user is entering st.subheader("Enter a word from which a sentence / word would be predicted") text2 = st.text_input('Enter word: ') #text is stored in this variable if option == '13M_OLD': option2 = st.selectbox('Type',('word','sentence')) if option2 == 'word': len2 = 1 else: len2 = 13 vocab_size = 100000 max_len = 1 d_model = 64 # 64 , 1024 n_head = 4 # 8 , 16 ff_dim = 256 # 256 , 2048 dropout_rate = 0.1 # 0.5 , 0.2 weights = "predict3" datafile = "data2.txt" dict = "dict_predict3.bin.npz" with open(datafile,"r") as f: text = f.read() text = text.lower() words = text.split() loaded_dict = np.load(dict, allow_pickle=True) word_to_num = loaded_dict["word_to_num"].item() num_to_word = loaded_dict["num_to_word"].item() X = [] Y = [] for i in range(len(words)-1): word = words[i] next_word = words[i+1] X.append(word_to_num[word]) Y.append(word_to_num[next_word]) Y.append(0) X.append(word_to_num[words[-1]]) X_train = pad_sequences([X]) y_train = pad_sequences([Y]) chatbot = TransformerChatbot(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate) chatbot.load_weights(weights) chatbot.build(input_shape=(None, max_len)) # Build the model chatbot.compile(optimizer="adam", loss="sparse_categorical_crossentropy") for i in range(1): other_text2 = text2 other_text2 = other_text2.lower() other_words2 = other_text2.split() other_num2 = [word_to_num[word] for word in other_words2] given_X2 = other_num2 input_sequence2 = pad_sequences([given_X2], maxlen=max_len, padding='post') output_sentence = other_text2 + "" for _ in range(len2): predicted_token = np.argmax(chatbot.predict(input_sequence2), axis=-1) predicted_token = predicted_token.item() out = num_to_word[predicted_token] # if out == ".": # break output_sentence += " " + out given_X2 = given_X2[1:] given_X2.append(predicted_token) input_sequence2 = pad_sequences([given_X2], maxlen=max_len, padding='post') out2 = output_sentence st.write("Predicted Text: ") st.write(out2) elif option=="26M_OLD": option2 = st.selectbox('Type',('word','sentence')) if option2 == 'word': len2 = 1 else: len2 = 13 vocab_size = 100000 max_len = 1 d_model = 128 # 64 , 1024 n_head = 4 # 8 , 16 ff_dim = 256 # 256 , 2048 dropout_rate = 0.1 # 0.5 , 0.2 weights = "predict1" datafile = "data2.txt" dict = "dict_predict1.bin.npz" vocab_size = 100000 max_len = 1 d_model = 64 # 64 , 1024 n_head = 4 # 8 , 16 ff_dim = 256 # 256 , 2048 dropout_rate = 0.1 # 0.5 , 0.2 weights = "predict3" datafile = "data2.txt" dict = "dict_predict3.bin.npz" with open(datafile,"r") as f: text = f.read() text = text.lower() words = text.split() loaded_dict = np.load(dict, allow_pickle=True) word_to_num = loaded_dict["word_to_num"].item() num_to_word = loaded_dict["num_to_word"].item() X = [] Y = [] for i in range(len(words)-1): word = words[i] next_word = words[i+1] X.append(word_to_num[word]) Y.append(word_to_num[next_word]) Y.append(0) X.append(word_to_num[words[-1]]) X_train = pad_sequences([X]) y_train = pad_sequences([Y]) chatbot = TransformerChatbot(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate) chatbot.load_weights(weights) chatbot.build(input_shape=(None, max_len)) # Build the model chatbot.compile(optimizer="adam", loss="sparse_categorical_crossentropy") for i in range(1): other_text2 = text2 other_text2 = other_text2.lower() other_words2 = other_text2.split() other_num2 = [word_to_num[word] for word in other_words2] given_X2 = other_num2 input_sequence2 = pad_sequences([given_X2], maxlen=max_len, padding='post') output_sentence = other_text2 + "" for _ in range(len2): predicted_token = np.argmax(chatbot.predict(input_sequence2), axis=-1) predicted_token = predicted_token.item() out = num_to_word[predicted_token] # if out == ".": # break output_sentence += " " + out given_X2 = given_X2[1:] given_X2.append(predicted_token) input_sequence2 = pad_sequences([given_X2], maxlen=max_len, padding='post') out2 = output_sentence st.write("Predicted Text: ") st.write(out2) else: out2 = "Error: Wrong Model Selected" st.write(out2)