app update

app.py

@@ -1,25 +1,62 @@
+# import tensorflow as tf
+# from google.protobuf.struct_pb2 import Struct
+# import struct2tensor.ops.gen_decode_proto_sparse
+
+# def create_proto_message(text):
+#     message = Struct()
+#     message.fields["task_data"].string_value = text
+#     return message.SerializeToString()
+
+# class TFProtoModel:
+#     def __init__(self, model_path):
+#         self.model = tf.saved_model.load(model_path)
+#         self.infer = self.model.signatures['serving_default']
+    
+#     def predict(self, text):
+#         proto_data = create_proto_message(text)
+#         input_tensor = tf.constant([proto_data], dtype=tf.string)
+#         result = self.infer(inputs=input_tensor)
+#         return result['outputs'].numpy()
+
+# # Initialize model when the file is loaded
+# model = TFProtoModel("model")
+
+# # This is the function Hugging Face will call
+# def pipeline(text):
+#     return model.predict(text)
+
+
 import tensorflow as tf
 from google.protobuf.struct_pb2 import Struct
+from transformers import Pipeline
+import struct2tensor.ops.gen_decode_proto_sparse
+
 
 def create_proto_message(text):
     message = Struct()
     message.fields["task_data"].string_value = text
     return message.SerializeToString()
 
-class TFProtoModel:
-    def __init__(self, model_path):
+class TFProtoModel(Pipeline):
+    def __init__(self, model_path="model"):
         self.model = tf.saved_model.load(model_path)
         self.infer = self.model.signatures['serving_default']
     
-    def predict(self, text):
+    def _sanitize_parameters(self, **kwargs):
+        return {}, {}, {}
+
+    def preprocess(self, text):
         proto_data = create_proto_message(text)
-        input_tensor = tf.constant([proto_data], dtype=tf.string)
+        return tf.constant([proto_data], dtype=tf.string)
+
+    def _forward(self, input_tensor):
         result = self.infer(inputs=input_tensor)
         return result['outputs'].numpy()
 
-# Initialize model when the file is loaded
-model = TFProtoModel("model")
+    def postprocess(self, model_outputs):
+        return {"score": float(model_outputs[0])}
+
+pipeline = TFProtoModel
 
-# This is the function Hugging Face will call
-def pipeline(text):
-    return model.predict(text)
+# To specify the task
+task = "text-classification"  # or another appropriate task type