File size: 2,609 Bytes
03b56f8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | # -*- coding: utf-8 -*-
"""
推論スクリプト
使い方:
python src/infer.py --model out/large "kyouhaiitenkidesune"
python src/infer.py --model out/large # 対話モード
"""
import argparse
import sys
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from normalization import normalize_input
BOS_IN = "\uEE00"
BOS_OUT = "\uEE01"
def convert(model, tok, romaji, device):
normalized = normalize_input(romaji)
prompt = BOS_IN + normalized + BOS_OUT
enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(device)
prompt_token_cap = int(enc.attention_mask.sum(dim=1).to("cpu").max().item()) + 32
generation_cap = min(prompt_token_cap, 768)
max_positions = getattr(model.config, "max_position_embeddings", None)
if max_positions:
remaining_positions = max_positions - enc.input_ids.shape[1]
if remaining_positions > 0:
generation_cap = min(generation_cap, remaining_positions)
else:
generation_cap = 1
generation_cap = max(1, generation_cap)
with torch.no_grad():
out = model.generate(
enc.input_ids,
attention_mask=enc.attention_mask,
max_new_tokens=generation_cap,
do_sample=False,
eos_token_id=tok.eos_token_id,
pad_token_id=tok.pad_token_id,
)
return tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)
def main():
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("text", nargs="?", default=None)
args = ap.parse_args()
device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(
args.model, dtype=torch.bfloat16 if device == "cuda" else torch.float32
).to(device).eval()
if args.text:
t0 = time.time()
print(convert(model, tok, args.text, device))
print(f"({(time.time()-t0)*1000:.0f} ms)")
else:
print("ローマ字を入力してください（空行で終了）")
while True:
try:
line = input("> ").strip()
except EOFError:
break
if not line:
break
t0 = time.time()
print(f" {convert(model, tok, line, device)} ({(time.time()-t0)*1000:.0f} ms)")
if __name__ == "__main__":
main()
|