fix: error when loading huggingface model with embedded language model

This commit is contained in:
Jonatas Grosman
2023-01-31 23:24:26 -03:00
parent acaefa09a1
commit d294e29ad9

View File

@ -5,7 +5,7 @@ C. Max Bain
import numpy as np import numpy as np
import pandas as pd import pandas as pd
from typing import List, Union, Iterator, TYPE_CHECKING from typing import List, Union, Iterator, TYPE_CHECKING
from transformers import AutoProcessor, Wav2Vec2ForCTC from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
import torchaudio import torchaudio
import torch import torch
from dataclasses import dataclass from dataclasses import dataclass
@ -52,7 +52,7 @@ def load_align_model(language_code, device, model_name=None):
align_dictionary = {c.lower(): i for i, c in enumerate(labels)} align_dictionary = {c.lower(): i for i, c in enumerate(labels)}
else: else:
try: try:
processor = AutoProcessor.from_pretrained(model_name) processor = Wav2Vec2Processor.from_pretrained(model_name)
align_model = Wav2Vec2ForCTC.from_pretrained(model_name) align_model = Wav2Vec2ForCTC.from_pretrained(model_name)
except Exception as e: except Exception as e:
print(e) print(e)