change handler to use whisper model and include segments

handler.py

@@ -2,15 +2,18 @@ from typing import  Dict
 from transformers.pipelines.audio_utils import ffmpeg_read
 import whisper
 import torch
-
-SAMPLE_RATE = 16000
-
+import pytube
 
 
 class EndpointHandler():
     def __init__(self, path=""):
         # load the model
-        self.model = whisper.load_model("medium")
+        MODEL_NAME = "tiny.en"
+        
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f'whisper will use: {device}')
+        
+        whisper_model = whisper.load_model(MODEL_NAME).to(device)
 
 
     def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:
@@ -19,15 +22,25 @@ class EndpointHandler():
             data (:obj:):
                 includes the deserialized audio file as bytes
         Return:
-            A :obj:`dict`:. base64 encoded image
+            A :obj:`dict`:. URL to image for download
         """
         # process input
-        inputs = data.pop("inputs", data)
-        audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
-        audio_tensor= torch.from_numpy(audio_nparray)
-        
-        # run inference pipeline
-        result = self.model.transcribe(audio_nparray)
+        print('data', data)
+        video_url = data.pop("inputs", data)
+        decode_options = {
+            # Set language to None to support multilingual, 
+            # but it will take longer to process while it detects the language.
+            # Realized this by running in verbose mode and seeing how much time
+            # was spent on the decoding language step
+            "language":"en",
+            verbose: True
+        }
+        yt = pt.YouTube(video_url)
+        stream = yt.streams.filter(only_audio=True)[0]
+        path_to_audio = f"{yt.video_id}.mp3"
+        stream.download(filename=path_to_audio)
+
+        transcript = self.model.transcribe(path_to_audio, **decode_options)
 
         # postprocess the prediction
-        return {"text": result["text"]}
+        return {"transcript": transcript}