File size: 1,409 Bytes
735386f 34a5c7b 735386f 34a5c7b |
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 |
from transformers import (
EncoderDecoderModel,
AutoTokenizer
)
import torch
PRETRAINED = "raynardj/wenyanwen-chinese-translate-to-ancient"
def inference(text):
tk_kwargs = dict(
truncation=True,
max_length=128,
padding="max_length",
return_tensors='pt')
inputs = tokenizer([text,],**tk_kwargs)
with torch.no_grad():
return tokenizer.batch_decode(
model.generate(
inputs.input_ids,
attention_mask=inputs.attention_mask,
num_beams=3,
bos_token_id=101,
eos_token_id=tokenizer.sep_token_id,
pad_token_id=tokenizer.pad_token_id,
), skip_special_tokens=True)
import streamlit as st
st.title("Wenyanwen Translator")
st.markdown("""
# Translate from Chinese to Ancient Chinese / 还你古朴清雅壮丽的文言文, 这[github](https://github.com/raynardj/yuan)
最多100个中文字符, 谢谢!
""")
@st.cache
def load_model():
tokenizer = AutoTokenizer.from_pretrained(PRETRAINED)
model = EncoderDecoderModel.from_pretrained(PRETRAINED)
return tokenizer, model
tokenizer, model = load_model()
text = st.text_area("轻轻地我走了,正如我轻轻地来。我挥一挥衣袖,不带走一片云彩。")
if st.button("Translate"):
if len(text) > 100:
st.error("超过100个字符")
else:
st.write(inference(text)[0])
|