swdq/ASR-Fixer

asrの文字起こし結果を用いて誤字を修正するために作成したモデルです。
主にデータセット整形に特化させたモデルです。
ルールベースではない誤字修正を行うことができることを目指しています。

使い方は下記コードを参照してください。

# 出力のスコアを含めて生成する
import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer


# CUDAデバイスの設定
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# モデルとトークナイザーのディレクトリ
MODEL_DIRECTORY = ("swdq/ASR-Fixer")


def load_model_and_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIRECTORY)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_DIRECTORY,
        device_map="auto",
        use_auth_token=True,
        torch_dtype=torch.float16,
    )
    return tokenizer, model


def generate_text(prompt, tokenizer, model):
    streamer = TextStreamer(tokenizer, skip_prompt=False)
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    # 出力のスコアを含めて生成する
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_new_tokens=500,
        min_new_tokens=1,
        do_sample=True,
        top_k=1,
        repetition_penalty=1.0,
        temperature=0.1,
        use_cache=True,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)


"""メインループ関数"""
tokenizer, model = load_model_and_tokenizer()

text = "れろれろ…じゅるる…このおちんちんなら…ちゅぱ、ちゅる…ちゃんとママの言うこと聞けまちゅよ…"
initial_prompt = f"ASR: {text} 正解: "
generated_text = generate_text(initial_prompt, tokenizer, model)
result = generated_text.split("正解: ", 1)[-1].strip()
print(text)
print(result)

例(データセットにないテキスト)
asr れろれろ…じゅるる…このおちんちんなら…ちゅぱ、ちゅる…ちゃんとママの言うこと聞けまちゅよ…
true れろれろ…ちゅるる…このおちんちんなら…ちゅぱ、ちゅる…ちゃんとママの言うこと、聞けまちゅよ…
asr 触られるとぴくっとする気持ちがわかりました くすぐったいような…気持ちいいような…んっ…
true 触られると、ぴくっとする気持ちがわかりました…くすぐったいような…気持ちいいような…んっ…
asr えへへ…私のお口で、ですけど… あはは…ってか、あは…やっぱり、結構勃っちゃってますよね、これ。
true えへへ…私のお口で、ですけど…あはは…ってか、あはは…やっぱり、結構勃っちゃってますよね、これ。
asr 坊ちゃま、おっ、おっ、お゛っ、おっ、おおっ、イグ…ッ、おぉ、おお゛お゛おぉぉッッ!
true 坊ちゃま、おっ、おっ、おぉっ、おおっ、イグ…ッ、おぉ、おぉ、おぉぉおおぉぉおおぉぉおおぉぉおおぉおおッッ!
asr さあさあ、おちんぽ出したんだから、次は便器に剥いて、おしっこシャーってしちゃお
true さあさあ、おちんぽ出したんだから、次は便器に剥いて、おしっこしゃーってしちゃお
asr 出して出して、ザーメン出して、ざーめん出して。
true 出して出して、ザーメン出して、ザーメン出して。
asr ざーめーん、せーえき 精子、まっしろを玉さくし…いっぱい眉の足にぶっかけてください。
true ざーめーん、せーえき

modelはLocal-Novel-LLM-project/Vecteus-v1をsftで独自データセットを用いて訓練しました。