Create app.py

app.py

@@ -0,0 +1,122 @@
+import tensorflow as tf
+from tensorflow.keras.layers import TextVectorization, Embedding, MultiHeadAttention, LayerNormalization, Dense, Dropout
+from tensorflow.keras.models import Model
+import gradio as gr
+import json
+
+START_TOKEN = '<start>'
+END_TOKEN = '<end>'
+
+class TransformerBlock(tf.keras.layers.Layer):
+    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.2, **kwargs):
+        super().__init__(**kwargs)
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.ff_dim = ff_dim
+        self.rate = rate
+
+        self.att = MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
+        self.ffn = tf.keras.Sequential([
+            Dense(ff_dim, activation='relu'),
+            Dense(embed_dim),
+        ])
+
+        self.layernorm1 = LayerNormalization(epsilon=1e-5)
+        self.layernorm2 = LayerNormalization(epsilon=1e-5)
+        self.dropout1 = Dropout(rate)
+        self.dropout2 = Dropout(rate)
+
+    def call(self, inputs, training=None):
+        attn_output = self.att(inputs, inputs)
+        attn_output = self.dropout1(attn_output, training=training)
+        out1 = self.layernorm1(inputs + attn_output)
+
+        ffn_output = self.ffn(out1)
+        ffn_output = self.dropout2(ffn_output, training=training)
+        return self.layernorm2(out1 + ffn_output)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            'embed_dim': self.embed_dim,
+            'num_heads': self.num_heads,
+            'ff_dim': self.ff_dim,
+            'rate': self.rate,
+        })
+        return config
+
+class TokenAndPositionEmbedding(tf.keras.layers.Layer):
+    def __init__(self, maxlen, vocab_size, embed_dim, **kwargs):
+        super().__init__(**kwargs)
+        self.maxlen = maxlen
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+
+        self.token_emb = Embedding(input_dim=vocab_size, output_dim=embed_dim)
+        self.pos_emb = Embedding(input_dim=maxlen, output_dim=embed_dim)
+
+    def call(self, x):
+        maxlen = tf.shape(x)[-1]
+        positions = tf.range(start=0, limit=maxlen, delta=1)
+        positions = self.pos_emb(positions)
+        x = self.token_emb(x)
+        return x + positions
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            'maxlen': self.maxlen,
+            'vocab_size': self.vocab_size,
+            'embed_dim': self.embed_dim,
+        })
+        return config
+
+def load_model(filename="tg-medium"):
+    model = tf.keras.models.load_model(f'{filename}.h5', custom_objects={
+        'TokenAndPositionEmbedding': TokenAndPositionEmbedding,
+        'TransformerBlock': TransformerBlock
+    })
+    with open(f'{filename}.json', 'r', encoding='utf-8') as f:
+        vocab = json.load(f)
+    vectorizer = TextVectorization(
+        max_tokens=96000,
+        output_sequence_length=100,
+        standardize=None,
+        vocabulary=vocab
+    )
+    return model, vectorizer
+
+def generate_text(model, vectorizer, prompt):
+    prompt = START_TOKEN + ' ' + prompt + ' ' + END_TOKEN
+    input_seq = vectorizer([prompt])
+    input_seq = input_seq[:, :-1]
+    predictions = model.predict(input_seq)
+    predicted_tokens = tf.argmax(predictions[0], axis=-1)
+    vocab = vectorizer.get_vocabulary()
+    output_tokens = [vocab[idx] for idx in predicted_tokens.numpy()]
+    if END_TOKEN in output_tokens:
+        end_index = output_tokens.index(END_TOKEN)
+        output_tokens = output_tokens[:end_index]
+    if START_TOKEN in output_tokens:
+        output_tokens.remove(START_TOKEN)
+    output = ' '.join(output_tokens)
+    return output
+
+def main():
+    model, vectorizer = load_model()
+
+    def generate_response(prompt):
+        return generate_text(model, vectorizer, prompt)
+
+    iface = gr.Interface(
+        fn=generate_response,
+        inputs=gr.Textbox(lines=2, placeholder="Start your conversation."),
+        outputs="text",
+        title="tg-medium",
+        description="Interference API. (russian only)"
+    )
+
+    iface.launch()
+
+if __name__ == "__main__":
+    main()