unitminer/urlslab_inference

1import os2from typing import Any3from sentence_transformers import SentenceTransformer4import pytube5from whisper_jax import FlaxWhisperPipline6import jax.numpy as jnp7import mimetypes8 9class EndpointHandler():10    def __init__(self, path=""):11        self.model =  SentenceTransformer('sentence-transformers/paraphrase-multilingual-mpnet-base-v2')12        self.pipeline = FlaxWhisperPipline("openai/whisper-tiny", dtype=jnp.float16, batch_size=16)13 14    def __call__(self, data: Any) -> Any:15        """16        Args:17            data (:obj:):18                includes the input data and the parameters for the inference.19        Return:20            A :obj:`list`:. The list contains the embeddings of the inference inputs21        """22 23        if data.get("task") == "transcribe":24 25            yt_id = data.get("inputs")26            if yt_id is None:27                return {"error": "inputs is required, should contain youtube id or youtube video url"}28            yt_url = 'https://youtu.be/' + yt_id29            filename = yt_id + ".m4a"30 31            # first try to load captions from YT32            yt = pytube.YouTube(yt_url)33            captions = yt.captions34            if captions:35                for c in captions:36                    caption = yt.captions.get_by_language_code(c.code)37                    return {"transcript": caption.generate_srt_captions()}38 39            if not os.path.isfile(filename):40 41                try:42 43                    stream = yt.streams.filter(only_audio=True)[0]44                    extension = mimetypes.guess_extension(stream.mime_type)45                    if extension:46                        filename = yt_id + extension47                    stream.download(filename=filename)48 49                except KeyError:50                    return {"error":"An error occurred while loading the YouTube video. Please try again."}51 52            try:53                if os.path.isfile(filename):54 55                    text = self.pipeline(filename, task="transcribe", return_timestamps=True)56 57                else:58                    text = "file not found"59            except Exception as e:60                return {"error": "An error occurred while transcribing the audio. Please try again."+ str(e)}61            finally:62                if os.path.isfile(filename):63                    os.remove(filename)64            return { "transcript": text}65        else:66            inputs = data.get("inputs")67            embeddings = self.model.encode(inputs)68            return {"embeddings": embeddings.tolist()}