1 develop transcription translation summarisation pipeline

pyproject.toml

@@ -31,7 +31,10 @@ dependencies = [
   "transformers",
   "huggingface",
   "datasets",
-  "numpy"
+  "numpy",
+  "sentencepiece",
+  "librosa",
+  "soundfile"
 ]
 
 [project.optional-dependencies]

scripts/TTS_pipeline_example.py

@@ -0,0 +1,46 @@
+"""
+    An example use of the transcription, translation and summarisation pipeline.
+"""
+
+import numpy as np
+from datasets import Audio, load_dataset
+
+from arc_spice.pipelines.TTS_pipeline import TTSpipeline
+
+
+def main(TTS_params):
+    """main function"""
+    TTS = TTSpipeline(TTS_params)
+    TTS.print_pipeline()
+    ds = load_dataset(
+        "facebook/multilingual_librispeech", "french", split="test", streaming=True
+    )
+    ds = ds.cast_column("audio", Audio(sampling_rate=16_000))
+    input_speech = next(iter(ds))["audio"]
+    # arrays = []
+    # n = 5
+    # for idx, data in enumerate(iter(ds)):
+    #     arrays.append(data["audio"]["array"])
+    #     if idx == n:
+    #         break
+    # arrays = np.concatenate(arrays)
+    TTS.run_pipeline(input_speech["array"])
+    TTS.print_results()
+
+
+if __name__ == "__main__":
+    TTS_pars = {
+        "transcriber": {
+            "specific_task": "automatic-speech-recognition",
+            "model": "openai/whisper-small",
+        },
+        "translator": {
+            "specific_task": "translation_fr_to_en",
+            "model": "facebook/mbart-large-50-many-to-many-mmt",
+        },
+        "summariser": {
+            "specific_task": "summarization",
+            "model": "facebook/bart-large-cnn",
+        },
+    }
+    main(TTS_params=TTS_pars)

src/arc_spice/pipelines/TTS_pipeline.py

@@ -0,0 +1,53 @@
+"""
+    Class for the transcription, translation and summarisation pipeline.
+"""
+
+from transformers import pipeline
+
+
+class TTSpipeline:
+    """
+    Class for the transcription, translation, summarisation pipeline.
+
+    pars:
+        - {'top_level_task': {'specific_task': str, 'model_name': str}}
+    """
+
+    def __init__(self, pars) -> None:
+        self.pars = pars
+        self.transcriber = pipeline(
+            pars["transcriber"]["specific_task"], pars["transcriber"]["model"]
+        )
+        self.translator = pipeline(
+            pars["translator"]["specific_task"], pars["translator"]["model"]
+        )
+        self.summariser = pipeline(
+            pars["summariser"]["specific_task"], pars["summariser"]["model"]
+        )
+        self.results = {}
+
+    def print_pipeline(self):
+        """Print the models in the pipeline"""
+        print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
+        print(f"Transcriber model: {self.pars['transcriber']['model']}")
+        print(f"Translator model: {self.pars['translator']['model']}")
+        print(f"Summariser model: {self.pars['summariser']['model']}")
+        print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
+
+    def run_pipeline(self, x):
+        """Run the pipeline on an input x"""
+        transcription = self.transcriber(x)
+        self.results["transcription"] = transcription["text"]
+        translation = self.translator(transcription["text"])
+        self.results["translation"] = translation[0]["translation_text"]
+        summarisation = self.summariser(translation[0]["translation_text"])
+        self.results["summarisation"] = summarisation[0]["summary_text"]
+
+    def print_results(self):
+        """Print the results for quick scanning"""
+        print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
+        for key, val in self.results.items():
+            print("-------------")
+            print(f"{key} result is: \n {val}")
+        print("-------------")
+        print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")