user model

configuration.py

@@ -0,0 +1,8 @@
+from transformers import PretrainedConfig
+
+
+class UserModelConfig(PretrainedConfig):
+    model_type = "user"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)

modelling.py

@@ -0,0 +1,14 @@
+from user_model.configuration import UserModelConfig
+from transformers import PreTrainedModel
+import tensorflow as tf
+
+
+class UserModel(PreTrainedModel):
+    config_class = UserModelConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = tf.saved_model.load('user_model/tf_retrieval_user_model')
+
+    def forward(self, user_id):
+        return self.model(user_id)

tf_retrieval_user_model/._saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1871400a8a58b60e5a850cf4516ede5bf8cb83f029fb5cd3138928df079d635
+size 581

tf_retrieval_user_model/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2348b5487a6c766d2a0ed904aefe4ea0ed77608973d68cc920344367ef53104e
+size 51237