File size: 1,352 Bytes
702c23d 09cc1b9 702c23d 09cc1b9 702c23d 09cc1b9 702c23d 09cc1b9 702c23d 09cc1b9 702c23d 09cc1b9 | 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 | import tika
tika.initVM()
from tika import parser
import pickle
import gradio as gr
from sklearn.pipeline import Pipeline
from tempfile import _TemporaryFileWrapper
doc_cls = [
"Договоры аренды",
"Договоры купли-продажи",
"Договоры оказания услуг",
"Договоры подряда",
"Договоры поставки"
]
class Classifier:
def __init__(self, pipeline: Pipeline):
self.pipeline = pipeline
def __call__(self, doc: _TemporaryFileWrapper):
if not doc:
return
doc.seek(0)
buffer = doc.read(-1)
parsed = parser.from_buffer(buffer)
content = parsed["content"]
probs = self.pipeline.predict_proba([content])[0]
return {d:p for d, p in zip(doc_cls, probs)}
def main():
tika.initVM()
with open("pipeline.pkl", "rb") as file:
pipeline: Pipeline = pickle.load(file)
classifier = Classifier(pipeline)
with gr.Blocks() as demo:
doc = gr.File(label="Документ")
output = gr.Label(label="Результаты классификации")
button = gr.Button(value="Классифицировать", variant="primary")
button.click(classifier, doc, output)
demo.launch()
if __name__ == "__main__":
main()
|