from transformers import AutoProcessor
from transformers import Wav2Vec2ProcessorWithLM
from huggingface_hub import Repository
from pyctcdecode import build_ctcdecoder

model_name = "cahya/wav2vec2-base-turkish-artificial-cv"
processor = AutoProcessor.from_pretrained(model_name)

vocab_dict = processor.tokenizer.get_vocab()
sorted_vocab_dict = {k.lower(): v for k, v in sorted(vocab_dict.items(), key=lambda item: item[1])}

decoder = build_ctcdecoder(
    labels=list(sorted_vocab_dict.keys()),
    kenlm_model_path="5gram.arpa",
)

processor_with_lm = Wav2Vec2ProcessorWithLM(
    feature_extractor=processor.feature_extractor,
    tokenizer=processor.tokenizer,
    decoder=decoder
)

#repo = Repository(local_dir="wav2vec2-base-turkish", clone_from=model_name)
processor_with_lm.save_pretrained("wav2vec2-base-turkish")
#repo.push_to_hub(commit_message="Upload lm-boosted decoder")