osanseviero HF staff commited on
Commit
b18f6c8
1 Parent(s): 2302c8e

Change docstring

Browse files
Files changed (1) hide show
  1. model.py +3 -2
model.py CHANGED
@@ -12,7 +12,7 @@ import os
12
  class PreTrainedModel(Runner):
13
  def __init__(self, path=""):
14
  """
15
- Loads model and tokenizer from local directory
16
  """
17
  ckp_file = os.path.join(path, "hubert_asr.ckpt")
18
  ckp = torch.load(ckp_file, map_location='cpu')
@@ -41,12 +41,13 @@ class PreTrainedModel(Runner):
41
  features = self.upstream.model(inputs)
42
  features = self.featurizer.model(inputs, features)
43
  preds = self.downstream.model.inference(features, [])
44
- return preds[0]
45
 
46
 
47
  """
48
  import subprocess
49
  import numpy as np
 
50
  # This is already done in the Inference API
51
  def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
52
  ar = f"{sampling_rate}"
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')
41
  features = self.upstream.model(inputs)
42
  features = self.featurizer.model(inputs, features)
43
  preds = self.downstream.model.inference(features, [])
44
+ return {"text": preds[0]}
45
 
46
 
47
  """
48
  import subprocess
49
  import numpy as np
50
+ from datasets import load_dataset
51
  # This is already done in the Inference API
52
  def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
53
  ar = f"{sampling_rate}"