1
"""
2
This is just an example of what people would submit for
3
inference.
4
"""
5
6
from s3prl.downstream.runner import Runner
7
from typing import Dict
8
import torch
9
import os
10
11
12
class PreTrainedModel(Runner):
13
    def __init__(self, path=""):
14
        """
15
        Initialize downstream model.
16
        """
17
        ckp_file = os.path.join(path, "hubert_asr.ckpt")
18
        ckp = torch.load(ckp_file, map_location='cpu')
19
        ckp["Args"].init_ckpt = ckp_file
20
        ckp["Args"].mode = "inference"
21
        ckp["Args"].upstream = "osanseviero/hubert_base"
22
        ckp["Args"].upstream_model_name = "model.pt"
23
        ckp["Args"].from_hf_hub = True
24
        ckp["Args"].device = "cpu" # Just to try in my computer
25
        ckp["Config"]["downstream_expert"]["datarc"]["dict_path"]=os.path.join(path,'char.dict')
26
27
        Runner.__init__(self, ckp["Args"], ckp["Config"])
28
29
    def __call__(self, inputs)-> Dict[str, str]:
30
        """
31
        Args:
32
            inputs (:obj:`np.array`):
33
                The raw waveform of audio received. By default at 16KHz.
34
        Return:
35
            A :obj:`dict`:. The object return should be liked {"text": "XXX"} containing
36
            the detected text from the input audio.
37
        """
38
        for entry in self.all_entries:
39
            entry.model.eval()
40
41
        inputs = [torch.FloatTensor(inputs)]
42
43
        with torch.no_grad():
44
            features = self.upstream.model(inputs)
45
            features = self.featurizer.model(inputs, features)
46
            preds = self.downstream.model.inference(features, [])
47
        return {"text": preds[0]}
48
49
"""
50
import subprocess
51
import numpy as np
52
from datasets import load_dataset
53
# This is already done in the Inference API
54
def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
55
    ar = f"{sampling_rate}"
56
    ac = "1"
57
    format_for_conversion = "f32le"
58
    ffmpeg_command = [
59
        "ffmpeg",
60
        "-i",
61
        "pipe:0",
62
        "-ac",
63
        ac,
64
        "-ar",
65
        ar,
66
        "-f",
67
        format_for_conversion,
68
        "-hide_banner",
69
        "-loglevel",
70
        "quiet",
71
        "pipe:1",
72
    ]
73
74
    ffmpeg_process = subprocess.Popen(
75
        ffmpeg_command, stdin=subprocess.PIPE, stdout=subprocess.PIPE
76
    )
77
    output_stream = ffmpeg_process.communicate(bpayload)
78
    out_bytes = output_stream[0]
79
80
    audio = np.frombuffer(out_bytes, np.float32).copy()
81
    if audio.shape[0] == 0:
82
        raise ValueError("Malformed soundfile")
83
    return audio
84
85
86
model = PreTrainedModel()
87
ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
88
filename = ds[0]["file"]
89
with open(filename, "rb") as f:
90
    data = ffmpeg_read(f.read(), 16000)
91
    print(model(data)) 
92
"""