|
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 |
|
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) |
|
|
|
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,len,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)) |
|
chatbot.compile(optimizer="adam", loss="sparse_categorical_crossentropy") |
|
|
|
for i in range(1): |
|
other_text1 = text2 |
|
other_text1 = other_text1.lower() |
|
other_words1 = other_text1.split() |
|
other_num1 = [word_to_num[word] for word in other_words1] |
|
given_X1 = other_num1 |
|
input_sequence1 = pad_sequences([given_X1], maxlen=max_len, padding='post') |
|
output_sentence = "" |
|
for _ in range(len): |
|
predicted_token = np.argmax(chatbot.predict(input_sequence1), axis=-1) |
|
predicted_token = predicted_token.item() |
|
out = num_to_word[predicted_token] |
|
|
|
|
|
output_sentence = out |
|
|
|
given_X1 = given_X1[1:] |
|
given_X1.append(predicted_token) |
|
input_sequence1 = pad_sequences([given_X1], 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") |
|
|
|
option = st.selectbox('Model',('13M','26M')) |
|
|
|
st.subheader("Enter a word from which a sentence / word would be predicted") |
|
|
|
text2 = st.text_input('Enter word: ') |
|
|
|
|
|
if option == '13M': |
|
option2 = st.selectbox('Type',('word','sentence')) |
|
if option2 == 'word': |
|
len = 1 |
|
else: |
|
len = 13 |
|
vocab_size = 100000 |
|
max_len = 1 |
|
d_model = 64 |
|
n_head = 4 |
|
ff_dim = 256 |
|
dropout_rate = 0.1 |
|
weights = "predict3" |
|
datafile = "data2.txt" |
|
dict = "dict_predict3.bin.npz" |
|
out2 = completion_model(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate,weights,datafile,dict,len,text2) |
|
st.write("Predicted Text: ") |
|
st.write(out2) |
|
|
|
|
|
elif option=="26M": |
|
option2 = st.selectbox('Type',('word','sentence')) |
|
if option2 == 'word': |
|
len = 1 |
|
else: |
|
len = 13 |
|
|
|
else: |
|
out2 = "Error: Wrong Model Selected" |
|
|
|
st.write(out2) |
|
|