Trend_Sentry / models.py
ItsProtesilaus's picture
Commit Alpha
b06572a
Raw
History Blame Contribute Delete
2.81 kB
import tensorflow as tf
import numpy as np
import tensorflow_hub as hub
from tensorflow.keras.utils import register_keras_serializable
@register_keras_serializable()
class USE(tf.keras.layers.Layer):
def __init__(self,
encoder_layer,
encoder_url="https://tfhub.dev/google/universal-sentence-encoder/4"
#encoder_url="https://www.kaggle.com/models/google/universal-sentence-encoder/frameworks/TensorFlow2/variations/universal-sentence-encoder"
, **kwargs):
super(USE, self).__init__(**kwargs)
self.encoder_layer=encoder_layer
self.encoder_url=encoder_url
def call(self, inputs):
return self.encoder_layer(inputs)
def get_config(self):
config=super(USE, self).get_config()
config['encoder_url']=self.encoder_url
return config
@classmethod
def from_config(cls, config):
encoder_url=config.pop('encoder_url')
encoder_layer=hub.KerasLayer(encoder_url, input_shape=[], dtype=tf.string, trainable=False, name='USE')
return cls(encoder_layer, encoder_url=encoder_url, **config)
encoder_layer=hub.KerasLayer("https://tfhub.dev/google/universal-sentence-encoder/4",
input_shape=[], ## The input is of variable length, hence the ip_length=[]
dtype=tf.string,
trainable=False,
name='USE')
class Model():
def __init__(self, encoder_layer=encoder_layer):
super(Model, self).__init__()
self.encoder_layer=encoder_layer
self.USE=USE(self.encoder_layer)
self.Sentiment_model=tf.keras.models.load_model('Best_sentiment_model.keras', custom_objects={'USE':self.USE})
self.Sentiment_model.trainable=False
self.Emotion_model=tf.keras.models.load_model('model_emotion_lstm.keras')
self.Emotion_model.trainable=False
def predict(self, text_sentiment=None, text_emotion=None):
sentiment = None
emotion = None
if text_sentiment is not None:
sentiment = self.Sentiment_model.predict(text_sentiment)
if text_emotion is not None:
emotion = self.Emotion_model.predict(text_emotion)
# If both were requested, return both
if text_sentiment is not None and text_emotion is not None:
return sentiment, emotion
# If only sentiment was requested
elif text_sentiment is not None:
return sentiment
# If only emotion was requested
elif text_emotion is not None:
return emotion
# If neither was provided
return None
if __name__=='__main__':
print(tf.__version__)
model=Model()
res1, res2=model.predict(tf.convert_to_tensor(['I am happy'], dtype=tf.string))
print(np.argmax(res1, axis=1), np.argmax(res2, axis=1))