| """Apply the punctuation/capitalisation tagger to a stream of ASR words. |
| |
| Words arrive one per line, lowercase and unpunctuated, exactly as the acoustic |
| model emits them. They are processed in overlapping windows so that a word near |
| a window edge is still decided with context on both sides; the middle of each |
| window wins. |
| """ |
| import argparse, json, sys, torch |
|
|
| PUNCT_LABELS = ["NONE", "COMMA", "PERIOD", "QUESTION", "EXCLAM", "COLON", |
| "SEMICOLON", "DASH"] |
| CASE_LABELS = ["LOWER", "TITLE", "UPPER"] |
| PUNCT_TEXT = {"NONE": "", "COMMA": ",", "PERIOD": ".", "QUESTION": "?", |
| "EXCLAM": "!", "COLON": ":", "SEMICOLON": ";", "DASH": " —"} |
| |
| TO_BASE = str.maketrans({"ҕ": "ғ", "ҥ": "ң"}) |
|
|
|
|
| class Tagger(torch.nn.Module): |
| def __init__(self, encoder, hidden): |
| super().__init__() |
| self.encoder = encoder |
| self.dropout = torch.nn.Dropout(0.1) |
| self.punct = torch.nn.Linear(hidden, len(PUNCT_LABELS)) |
| self.case = torch.nn.Linear(hidden, len(CASE_LABELS)) |
|
|
| def forward(self, input_ids, attention_mask): |
| h = self.encoder(input_ids=input_ids, |
| attention_mask=attention_mask).last_hidden_state |
| h = self.dropout(h) |
| return self.punct(h), self.case(h) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--model_dir", required=True) |
| p.add_argument("--words", required=True) |
| p.add_argument("--out", required=True) |
| p.add_argument("--window", type=int, default=120) |
| p.add_argument("--stride", type=int, default=80) |
| p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| args = p.parse_args() |
|
|
| from transformers import AutoModel, AutoTokenizer |
| ck = torch.load(args.model_dir + "/punct_model.pt", map_location="cpu", |
| weights_only=False) |
| tok = AutoTokenizer.from_pretrained(ck["base"]) |
| enc = AutoModel.from_pretrained(ck["base"]) |
| model = Tagger(enc, enc.config.hidden_size) |
| model.load_state_dict(ck["state_dict"]) |
| model.eval().to(args.device) |
|
|
| words = [w.strip() for w in open(args.words, encoding="utf-8") if w.strip()] |
| print("words: %d" % len(words), flush=True) |
|
|
| votes = [dict() for _ in words] |
| for start in range(0, len(words), args.stride): |
| chunk = words[start:start + args.window] |
| if not chunk: |
| break |
| e = tok([w.translate(TO_BASE) for w in chunk], is_split_into_words=True, |
| truncation=True, max_length=512, return_tensors="pt") |
| wid = e.word_ids() |
| with torch.inference_mode(): |
| pl, cl = model(e["input_ids"].to(args.device), |
| e["attention_mask"].to(args.device)) |
| pp, cc = pl.argmax(-1)[0].cpu(), cl.argmax(-1)[0].cpu() |
| prev = None |
| for t, w in enumerate(wid): |
| if w is None or w == prev: |
| prev = w |
| continue |
| prev = w |
| i = start + w |
| if i >= len(words): |
| continue |
| |
| d = abs(w - len(chunk) / 2) |
| if "d" not in votes[i] or d < votes[i]["d"]: |
| votes[i] = {"d": d, "punct": PUNCT_LABELS[pp[t]], |
| "case": CASE_LABELS[cc[t]]} |
| if start + args.window >= len(words): |
| break |
|
|
| out = [] |
| for w, v in zip(words, votes): |
| out.append({"word": w, "punct": v.get("punct", "NONE"), |
| "case": v.get("case", "LOWER")}) |
| json.dump(out, open(args.out, "w"), ensure_ascii=False) |
|
|
| def render(o): |
| s = o["word"] |
| if o["case"] == "TITLE": |
| s = s[:1].upper() + s[1:] |
| elif o["case"] == "UPPER": |
| s = s.upper() |
| return s + PUNCT_TEXT[o["punct"]] |
|
|
| print(" ".join(render(o) for o in out[:60]), flush=True) |
| print("\nwrote %s" % args.out) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
