| 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() |
|
|
