raynardj's picture
🎬 inference cache
89a5f1b
raw
history blame
1.5 kB
from transformers import (
EncoderDecoderModel,
AutoTokenizer
)
import torch
import streamlit as st
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)[0].replace(" ","")
st.title("🪕古朴 ❄️清雅 🌊壮丽")
st.markdown("""
> Translate from Chinese to Ancient Chinese / 还你古朴清雅壮丽的文言文, 这[github](https://github.com/raynardj/yuan)
> 最多100个中文字符
""")
@st.cache(allow_output_mutation=True)
def load_model():
tokenizer = AutoTokenizer.from_pretrained(PRETRAINED)
model = EncoderDecoderModel.from_pretrained(PRETRAINED)
return tokenizer, model
tokenizer, model = load_model()
text = st.text_area(value="轻轻地我走了,正如我轻轻地来。我挥一挥衣袖,不带走一片云彩。", label="输入文本")
if st.button("曰"):
if len(text) > 100:
st.error("无过百字,若过则当答此言。")
else:
st.write(inference(text))