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