from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import torch
import gradio as gr
import re

model_path = 'larryvrh/mt5-translation-ja_zh'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path, torch_dtype = torch.float32)


def paragraph_to_sentences(paragraph):
    sentences = []
    cursor = 0
    for i, c in enumerate(paragraph):
        if c == '。':
            sentences.append(paragraph[cursor:i + 1])
            cursor = i + 1
    if paragraph[-1] != '。':
        sentences.append(paragraph[cursor:])
    return sentences


def article_to_sentences(article):
    paragraphs = re.split(r'([\r\n\t ]+)', article)
    sentences = []
    for i, p in enumerate(paragraphs):
        if len(p.strip()) == 0:
            sentences.append(p)
        else:
            sentences += paragraph_to_sentences(p)
    return sentences

def stream_generate(model, tokenizer, *args, **kwargs):
    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
    kwargs['streamer'] = streamer
    t = Thread(target=model.generate, args=args, kwargs=kwargs)
    t.start()
    for c in streamer:
        yield c

def contains_CJK(text):
    return len(re.findall(r'[\u3040-\u309F\u30A0-\u30FF\uAC00-\uD7A3\u4E00-\u9FFF]', text)) > 0

def translate(text):
    result = ''
    for s in article_to_sentences(text):
        if not contains_CJK(s):
            result += s
            yield result
            continue
        for t in stream_generate(model, tokenizer, **tokenizer(f'<-ja2zh-> {s}', return_tensors="pt"),
                                 max_new_tokens=256, repetition_penalty=1.1, num_beams=1):
            result += t.replace(',','，').replace('!','！').replace('?','？').replace(';', '；').replace('(','（').replace(')','）')
            yield result

app = gr.Interface(
    fn=translate,
    inputs=gr.Textbox(lines=5),
    outputs=gr.Textbox(lines=5),
    allow_flagging='never',
    title='MT5-Translation-Ja_Zh',
    examples=[
        '文は、「主語・修飾語・述語」の語順で構成される。修飾語は被修飾語の前に位置する。また、名詞の格を示すためには、語順や語尾を変化させるのでなく、文法的な機能を示す機能語（助詞）を後ろに付け加える（膠着させる）。これらのことから、言語類型論上は、語順の点ではSOV型の言語に、形態の点では膠着語に分類される（「文法」の節参照）。',
        'フェルディナント・ラッサールは、プロイセンの政治学者、哲学者、法学者、社会主義者、労働運動指導者。ドイツ社会民主党の母体となる全ドイツ労働者同盟の創設者である。社会主義共和政の統一ドイツを目指しつつも、……',
        '広辞苑第七版によれば、百科事典は「学術・技芸・社会・家庭その他あらゆる科目にわたる知識を集め記し、これを部門別あるいは五十音順などに配列し、解説を加えた書物」のことであると定義しており、大辞泉では「人類の知識の及ぶあらゆる分野の事柄について、辞書の形式に準じて項目を立てて配列し、解説を加えた書物」であると定義されている。 ',
        'ウィキペディアは、信頼されるフリーな百科事典を、それも、質においても量においても史上最高の百科事典を目指して、共同作業で創り上げるプロジェクトです。あなたは、ウィキペディアをご覧になるだけでなく、今すぐにでも記事の編集に参加することができます。特別な参加資格はありません。編集作業には、あなたのパソコンのウェブブラウザ以外には特別な道具は必要ありません。これまでにない人類の知的遺産を育むこの壮大なプロジェクトにぜひ参加してください。あなたが関心のある分野、得意とする分野において、あなたの力を貸してください。 '
        ]
)
app.launch(enable_queue=True, max_threads=4)