Add ability to switch output languages for multilingual models #69

examples/speech_to_speech/reference/tgt_lang.txt

@@ -0,0 +1 @@
+es

examples/speech_to_text/counter_in_tgt_lang_agent.py

@@ -14,14 +14,10 @@ class CounterInTargetLanguage(SpeechToTextAgent):
     def __init__(self, args):
         super().__init__(args)
         self.wait_seconds = args.wait_seconds
-        self.tgt_lang = args.tgt_lang
 
     @staticmethod
     def add_args(parser):
         parser.add_argument("--wait-seconds", default=1, type=int)
-        parser.add_argument(
-            "--tgt-lang", default="en", type=str, choices=["en", "es", "de"]
-        )
 
     def policy(self, states: Optional[AgentStates] = None):
         if states is None:
@@ -35,11 +31,12 @@ def policy(self, states: Optional[AgentStates] = None):
             return ReadAction()
 
         prediction = f"{length_in_seconds} "
-        if self.tgt_lang == "en":
+        tgt_lang = states.tgt_lang
+        if tgt_lang == "en":
             prediction += "seconds"
-        elif self.tgt_lang == "es":
+        elif tgt_lang == "es":
             prediction += "segundos"
-        elif self.tgt_lang == "de":
+        elif tgt_lang == "de":
             prediction += "sekunden"
         else:
             prediction += "<unknown>"

examples/speech_to_text/eval.sh

@@ -1,5 +1,7 @@
 simuleval \
-    --agent english_counter_agent.py \
+    --agent counter_in_tgt_lang_agent.py \
     --source-segment-size 1000 \
     --source source.txt --target reference/en.txt \
+    --tgt-lang reference/tgt_lang.txt \
     --output output 
+

examples/speech_to_text/reference/tgt_lang.txt

@@ -0,0 +1 @@
+es

simuleval/agents/agent.py

@@ -196,6 +196,7 @@ class SpeechToTextAgent(GenericAgent):
 
     source_type: str = "speech"
     target_type: str = "text"
+    tgt_lang: Optional[str] = None
 
 
 class SpeechToSpeechAgent(GenericAgent):

simuleval/agents/states.py

@@ -29,6 +29,7 @@ def reset(self) -> None:
         self.target_finished = False
         self.source_sample_rate = 0
         self.target_sample_rate = 0
+        self.tgt_lang = None
         self.upstream_states = []
 
     def update_source(self, segment: Segment):
@@ -43,9 +44,11 @@ def update_source(self, segment: Segment):
             return
         elif isinstance(segment, TextSegment):
             self.source.append(segment.content)
+            self.tgt_lang = segment.tgt_lang
         elif isinstance(segment, SpeechSegment):
             self.source += segment.content
             self.source_sample_rate = segment.sample_rate
+            self.tgt_lang = segment.tgt_lang
         else:
             raise NotImplementedError
 

simuleval/data/dataloader/__init__.py

@@ -5,6 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
+from argparse import Namespace
+
 from .dataloader import (  # noqa
     GenericDataloader,
     register_dataloader,
@@ -21,7 +23,7 @@
 logger = logging.getLogger("simuleval.dataloader")
 
 
-def build_dataloader(args) -> GenericDataloader:
+def build_dataloader(args: Namespace) -> GenericDataloader:
     dataloader_key = getattr(args, "dataloader", None)
     if dataloader_key is not None:
         assert dataloader_key in DATALOADER_DICT, f"{dataloader_key} is not defined"

simuleval/data/dataloader/dataloader.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Optional
 from argparse import Namespace, ArgumentParser
 
 SUPPORTED_MEDIUM = ["text", "speech"]
@@ -37,10 +37,14 @@ class GenericDataloader:
     """
 
     def __init__(
-        self, source_list: List[str], target_list: Union[List[str], List[None]]
+        self,
+        source_list: List[str],
+        target_list: Union[List[str], List[None]],
+        tgt_lang_list: Optional[List[str]] = None,
     ) -> None:
         self.source_list = source_list
         self.target_list = target_list
+        self.tgt_lang_list = tgt_lang_list
         assert len(self.source_list) == len(self.target_list)
 
     def __len__(self):
@@ -52,8 +56,18 @@ def get_source(self, index: int) -> Any:
     def get_target(self, index: int) -> Any:
         return self.preprocess_target(self.target_list[index])
 
+    def get_tgt_lang(self, index: int) -> Optional[str]:
+        if self.tgt_lang_list is None or index >= len(self.tgt_lang_list):
+            return None
+        else:
+            return self.tgt_lang_list[index]
+
     def __getitem__(self, index: int) -> Dict[str, Any]:
-        return {"source": self.get_source(index), "target": self.get_target(index)}
+        return {
+            "source": self.get_source(index),
+            "target": self.get_target(index),
+            "tgt_lang": self.get_tgt_lang(index),
+        }
 
     def preprocess_source(self, source: Any) -> Any:
         raise NotImplementedError
@@ -95,3 +109,9 @@ def add_args(parser: ArgumentParser):
             default=1,
             help="Source segment size, For text the unit is # token, for speech is ms",
         )
+        parser.add_argument(
+            "--tgt-lang",
+            type=str,
+            default=None,
+            help="Target language",
+        )

simuleval/data/dataloader/s2t_dataloader.py

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 from pathlib import Path
-from typing import List, Union
+from typing import List, Union, Optional
 from .dataloader import GenericDataloader
 from simuleval.data.dataloader import register_dataloader
 from argparse import Namespace
@@ -58,6 +58,14 @@ def get_video_id(url):
 
 @register_dataloader("speech-to-text")
 class SpeechToTextDataloader(GenericDataloader):
+    def __init__(
+        self,
+        source_list: List[str],
+        target_list: List[str],
+        tgt_lang_list: Optional[List[str]] = None,
+    ) -> None:
+        super().__init__(source_list, target_list, tgt_lang_list)
+
     def preprocess_source(self, source: Union[Path, str]) -> List[float]:
         assert IS_IMPORT_SOUNDFILE, "Please make sure soundfile is properly installed."
         samples, _ = soundfile.read(source, dtype="float32")
@@ -75,40 +83,50 @@ def get_source_audio_path(self, index: int):
 
     @classmethod
     def from_files(
-        cls, source: Union[Path, str], target: Union[Path, str]
+        cls,
+        source: Union[Path, str],
+        target: Union[Path, str],
+        tgt_lang: Union[Path, str],
     ) -> SpeechToTextDataloader:
         with open(source) as f:
             source_list = [line.strip() for line in f]
         with open(target) as f:
             target_list = [line.strip() for line in f]
-        dataloader = cls(source_list, target_list)
+        with open(tgt_lang) as f:
+            tgt_lang_list = [line.strip() for line in f]
+        dataloader = cls(source_list, target_list, tgt_lang_list)
         return dataloader
 
     @classmethod
     def from_args(cls, args: Namespace):
         args.source_type = "speech"
         args.target_type = "text"
-        return cls.from_files(args.source, args.target)
+        return cls.from_files(args.source, args.target, args.tgt_lang)
 
 
 @register_dataloader("speech-to-speech")
 class SpeechToSpeechDataloader(SpeechToTextDataloader):
     @classmethod
     def from_files(
-        cls, source: Union[Path, str], target: Union[Path, str]
+        cls,
+        source: Union[Path, str],
+        target: Union[Path, str],
+        tgt_lang: Union[Path, str],
     ) -> SpeechToSpeechDataloader:
         with open(source) as f:
             source_list = [line.strip() for line in f]
         with open(target) as f:
             target_list = [line.strip() for line in f]
-        dataloader = cls(source_list, target_list)
+        with open(tgt_lang, "r") as f:
+            tgt_lang_list = [line.strip() for line in f]
+        dataloader = cls(source_list, target_list, tgt_lang_list)
         return dataloader
 
     @classmethod
     def from_args(cls, args: Namespace):
         args.source_type = "speech"
         args.target_type = "speech"
-        return cls.from_files(args.source, args.target)
+        return cls.from_files(args.source, args.target, args.tgt_lang)
 
 
 @register_dataloader("youtube-to-text")

simuleval/data/segments.py

@@ -6,6 +6,7 @@
 
 import json
 from dataclasses import dataclass, field
+from typing import Optional
 
 
 @dataclass
@@ -15,6 +16,7 @@ class Segment:
     finished: bool = False
     is_empty: bool = False
     data_type: str = None
+    tgt_lang: str = None
 
     def json(self) -> str:
         info_dict = {attribute: value for attribute, value in self.__dict__.items()}
@@ -34,12 +36,14 @@ class EmptySegment(Segment):
 class TextSegment(Segment):
     content: str = ""
     data_type: str = "text"
+    tgt_lang: str = Optional[str]
 
 
 @dataclass
 class SpeechSegment(Segment):
     sample_rate: int = -1
     data_type: str = "speech"
+    tgt_lang: str = Optional[str]
 
 
 def segment_from_json_string(string: str):

simuleval/evaluator/instance.py

@@ -46,6 +46,14 @@ def __init__(
         if self.dataloader is not None:
             self.source = self.dataloader[self.index]["source"]
             self.reference = self.dataloader[self.index]["target"]
+
+            if self.dataloader.tgt_lang_list is not None and self.index < len(
+                self.dataloader.tgt_lang_list
+            ):
+                self.tgt_lang = self.dataloader[self.index]["tgt_lang"]
+            else:
+                self.tgt_lang = None
+
         self.reset()
         if args is not None:
             self.args = args
@@ -235,6 +243,7 @@ def __init__(
         args: Optional[Namespace],
     ):
         super().__init__(index, dataloader, args)
+        self.args = args
         self.sample_rate_value = None
         self.sample_list = None
         self.source_finished_reading = False
@@ -282,6 +291,7 @@ def send_source(self, segment_size=10):
                 content=samples,
                 sample_rate=self.audio_info.samplerate,
                 finished=is_finished,
+                tgt_lang=self.tgt_lang,
             )
 
         else:

simuleval/options.py

@@ -161,6 +161,12 @@ def general_parser():
         default="main.yaml",
         help="Name of the config yaml of the system configs.",
     )
+    # parser.add_argument(
+    #     "--tgt-lang",
+    #     type=str,
+    #     default=None,
+    #     help="Path to the Target language file.",
+    # )
     parser.add_argument("--dataloader", default=None, help="Dataloader to use")
     parser.add_argument(
         "--log-level",

simuleval/test/test_s2s.py

@@ -39,6 +39,10 @@ def test_s2s(root_path=ROOT_PATH):
             os.path.join(root_path, "examples", "speech_to_speech", "reference/en.txt"),
             "--output",
             tmpdirname,
+            "--tgt-lang",
+            os.path.join(
+                root_path, "examples", "speech_to_speech", "reference/tgt_lang.txt"
+            ),
         ]
         cli.main()
 

simuleval/test/test_s2t.py

@@ -24,7 +24,7 @@ def test_s2t(root_path=ROOT_PATH):
         cli.sys.argv[1:] = [
             "--agent",
             os.path.join(
-                root_path, "examples", "speech_to_text", "english_counter_agent.py"
+                root_path, "examples", "speech_to_text", "counter_in_tgt_lang_agent.py"
             ),
             "--user-dir",
             os.path.join(root_path, "examples"),
@@ -38,6 +38,10 @@ def test_s2t(root_path=ROOT_PATH):
             os.path.join(root_path, "examples", "speech_to_text", "reference/en.txt"),
             "--output",
             tmpdirname,
+            "--tgt-lang",
+            os.path.join(
+                root_path, "examples", "speech_to_text", "reference/tgt_lang.txt"
+            ),
         ]
         cli.main()
 
@@ -46,7 +50,7 @@ def test_s2t(root_path=ROOT_PATH):
                 instance = LogInstance(line.strip())
                 assert (
                     instance.prediction
-                    == "1 second 2 second 3 second 4 second 5 second 6 second 7 second"
+                    == "1 segundos 2 segundos 3 segundos 4 segundos 5 segundos 6 segundos 7 segundos"
                 )
 
 
@@ -103,7 +107,9 @@ def test_s2t_with_tgt_lang(root_path=ROOT_PATH):
             "--output",
             tmpdirname,
             "--tgt-lang",
-            "es",
+            os.path.join(
+                root_path, "examples", "speech_to_text", "reference/tgt_lang.txt"
+            ),
         ]
         cli.main()