Update app.py

app.py

@@ -39,20 +39,12 @@ class TransformerChatbot(Model):
     def create_padding_mask(self, seq):
         mask = tf.cast(tf.math.equal(seq, 0), tf.float32)
         return mask[:, tf.newaxis, tf.newaxis, :]
-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',('1','2')) #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 == '1':
-    with open("data2.txt","r") as f:
+def textcompletion_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_predict3.bin.npz", allow_pickle=True)
+    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 = []
@@ -67,16 +59,11 @@ if option == '1':
     X.append(word_to_num[words[-1]])
     X_train = pad_sequences([X])
     y_train = pad_sequences([Y])
-    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
+    
 
 
     chatbot = TransformerChatbot(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate)
-    chatbot.load_weights("predict3")
+    chatbot.load_weights(weights)
     chatbot.build(input_shape=(None, max_len)) # Build the model
     chatbot.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
     
@@ -88,7 +75,7 @@ if option == '1':
         given_X1 = other_num1
         input_sequence1 = pad_sequences([given_X1], maxlen=max_len, padding='post')
         output_sentence = ""
-        for _ in range(1):
+        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]
@@ -101,63 +88,45 @@ if option == '1':
             input_sequence1 = pad_sequences([given_X1], maxlen=max_len, padding='post')
 
         out2 = output_sentence
-        
-    
-else:
-    with open("data2.txt","r") as f:
-        text = f.read()
-    text = text.lower()
-    words = text.split()
-    loaded_dict = np.load("dict_predict3.bin.npz", 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)
+        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','26M')) #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")
+len2 = st.text_input('Enter sequence length: ') 
+text2 = st.text_input('Enter word: ') #text is stored in this variable
 
-    X.append(word_to_num[words[-1]])
-    X_train = pad_sequences([X])
-    y_train = pad_sequences([Y])
+
+if option == '13M':
     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
-
-
-    chatbot = TransformerChatbot(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate)
-    chatbot.load_weights("predict3")
-    chatbot.build(input_shape=(None, max_len)) # Build the model
-    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 = other_text1+""
-        for _ in range(10):
-            predicted_token = np.argmax(chatbot.predict(input_sequence1), axis=-1)
-            predicted_token = predicted_token.item()
-            out = num_to_word[predicted_token]
-            
+    weights = "predict3"
+    datafile = "data2.txt"
+    dict = "dict_predict3.bin.npz"
+    len = len2
+    text2 = text2
+    out2 = textcompletion_model(vocab_size, max_len, d_model, n_head, ff_dim, dropout_rate,weights, datafile , dict, len , text2)
+        
     
-            output_sentence += " " + out
-            if out == ".":
-                break
-            given_X1 = given_X1[1:]
-            given_X1.append(predicted_token)
-            input_sequence1 = pad_sequences([given_X1], maxlen=max_len, padding='post')
-
-        out2 = output_sentence
+elif option=="26M":
+    vocab_size = 100000
+    max_len = 1
+    d_model = 128  # 64 , 1024
+    n_head = 8 # 8 , 16
+    ff_dim = 256  # 256 , 2048
+    dropout_rate = 0.1  # 0.5 , 0.2
+    weights = "predict5"
+    datafile = "data2.txt"
+    dict = "dict_predict3.bin.npz"
+    len = len2
+    text2 = text2
+    out2 = textcompletion_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)