# -*- coding: utf-8 -*-

import torch
import gradio as gr

from sbert_punc_case_ru import SbertPuncCase


device = "cuda" if torch.cuda.is_available() else "cpu"
sbertpunc = SbertPuncCase().to(device)


def punctuate(input_string):
    return sbertpunc.punctuate(input_string)


demo = gr.Interface(
    fn=punctuate,
    inputs="text",
    outputs="text",
    examples=["sbert punc case расставляет точки запятые и знаки вопроса вам нравится"],
    title="SbertPuncCase",
    description="SbertPuncCase - модель восстановления пунктуации и регистра для русского языка. Модель способна расставлять точки, запятые и знаки вопроса; определять регистр - слово в нижнем регистре, слово с первой буквой в верхнем регистре, слово в верхнем регистре. Модель разработана для восстановления текста после распознавания речи, поэтому работает со строками в нижнем регистре. В основу модели легла sbert_large_nlu_ru. В качестве обучающих данных использованы текстовые расшифровки интервью.",
    allow_flagging="never",
)
demo.launch()