Upload feature_extraction_moment.py

feature_extraction_moment.py

@@ -23,7 +23,7 @@ class MomentFeatureExtractor(FeatureExtractionMixin):
     # TODO: 本来はMoment側のTokenizerもts_tokenizerとして入れたかったが、モデルに組み込まれてしまっている。
     # refers: https://github.com/moment-timeseries-foundation-model/moment/blob/088b253a1138ac7e48a7efc9bf902336c9eec8d9/momentfm/models/moment.py#L105
 
-    model_input_names = ["time_series_values", "input_mask"]
+    model_input_names = ["time_series_values", "time_series_input_mask"]
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -50,7 +50,7 @@ class MomentFeatureExtractor(FeatureExtractionMixin):
             time_series_values = None
             input_mask = None
 
-        return BatchFeature(data={"time_series_values": time_series_values, "input_mask": input_mask})
+        return BatchFeature(data={"time_series_values": time_series_values, "time_series_input_mask": input_mask})
 
 
     def _convert_time_series(self, time_series, return_tensors, torch_dtype, padding, max_length):