/ test / python / testpipeline / testaudio / testtranscription.py
testtranscription.py
  1  """
  2  Transcription module tests
  3  """
  4  
  5  import unittest
  6  
  7  import numpy as np
  8  import soundfile as sf
  9  
 10  from scipy import signal
 11  
 12  from txtai.pipeline import Transcription
 13  
 14  # pylint: disable=C0411
 15  from utils import Utils
 16  
 17  
 18  class TestTranscription(unittest.TestCase):
 19      """
 20      Transcription tests.
 21      """
 22  
 23      def testArray(self):
 24          """
 25          Test audio data to text transcription
 26          """
 27  
 28          transcribe = Transcription()
 29  
 30          # Read audio data
 31          raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
 32  
 33          self.assertEqual(transcribe((raw, samplerate)), "Make huge profits without working make up to one hundred thousand dollars a day")
 34          self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")
 35  
 36      def testChunks(self):
 37          """
 38          Test splitting transcription into chunks
 39          """
 40  
 41          transcribe = Transcription()
 42  
 43          result = transcribe(Utils.PATH + "/Make_huge_profits.wav", join=False)[0]
 44  
 45          self.assertIsInstance(result["raw"], np.ndarray)
 46          self.assertIsNotNone(result["rate"])
 47          self.assertEqual(result["text"], "Make huge profits without working make up to one hundred thousand dollars a day")
 48  
 49      def testFile(self):
 50          """
 51          Test audio file to text transcription
 52          """
 53  
 54          transcribe = Transcription()
 55  
 56          self.assertEqual(
 57              transcribe(Utils.PATH + "/Make_huge_profits.wav"), "Make huge profits without working make up to one hundred thousand dollars a day"
 58          )
 59  
 60      def testGenerateArguments(self):
 61          """
 62          Test transcription with generation keyword arguments
 63          """
 64  
 65          transcribe = Transcription()
 66  
 67          # Read audio data
 68          raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
 69  
 70          self.assertEqual(
 71              transcribe(raw, samplerate, language="English", task="transcribe"),
 72              "Make huge profits without working make up to one hundred thousand dollars a day",
 73          )
 74  
 75      def testResample(self):
 76          """
 77          Test resampled audio file to text transcription
 78          """
 79  
 80          transcribe = Transcription()
 81  
 82          # Read audio data
 83          raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
 84  
 85          # Resample for testing
 86          samples = round(len(raw) * float(22050) / samplerate)
 87          raw, samplerate = signal.resample(raw, samples), 22050
 88  
 89          self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")
 90  
 91      def testStereo(self):
 92          """
 93          Test audio file in stereo to text transcription
 94          """
 95  
 96          transcribe = Transcription()
 97  
 98          # Read audio data
 99          raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav")
100  
101          # Convert mono to stereo
102          raw = np.column_stack((raw, raw))
103  
104          self.assertEqual(transcribe(raw, samplerate), "Make huge profits without working make up to one hundred thousand dollars a day")