Upload seamless_communication/datasets/datatypes.py with huggingface_hub

seamless_communication/datasets/datatypes.py

@@ -0,0 +1,47 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# MIT_LICENSE file in the root directory of this source tree.
+
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import torch
+
+
+@dataclass
+class MultimodalSample:
+    id: int
+    lang: str
+    text: str
+    audio_local_path: Optional[str] = None
+    waveform: Optional[torch.Tensor] = None
+    sampling_rate: Optional[int] = None
+    units: Optional[List[int]] = None
+
+    @classmethod
+    def from_json(cls, js: Dict[str, Any]) -> "MultimodalSample":
+        return cls(
+            id=js["id"],
+            lang=js["lang"],
+            text=js["text"],
+            audio_local_path=js.get("audio_local_path"),
+            waveform=None,  # don't serialize
+            sampling_rate=js.get("sampling_rate"),
+            units=js.get("units"),
+        )
+
+
+@dataclass
+class LangPairSample:
+    source: MultimodalSample
+    target: MultimodalSample
+
+    @classmethod
+    def from_json(cls, js: Dict[str, Any]) -> "LangPairSample":
+        return cls(
+            source=MultimodalSample.from_json(js["source"]),
+            target=MultimodalSample.from_json(js["target"]),
+        )