app.py

import hashlib
import logging
import os
from datetime import datetime
from typing import Dict, Any, List

import gradio as gr
import spacy
from pony.orm import db_session
from spacy import displacy

from core.classificator import Classificator, Classification
from core.models.request import Request
from settings import db


class Application:
    examples = [
        'Things are complicated because we still live together but we have separate lives',
        'My dad is a monster and took his anger out on my mom by verbally abusing her and when she left he '
        'eventually moved on to my brother',
        'A two months ago, she was chatting with some random guy',
        'Not I have a horrid relationship with my brother we’ve never gotten along and probably never will',
        'I was outside trying to leave and he caught me to explain why Im so rude',
    ]

    def __init__(self, classificator: Classificator, options: Dict[str, Any]):
        self.options = options
        self.classificator = classificator
        self.nlp = spacy.load("en_core_web_md")

    def handle(self, input_text: str) -> str:
        """
            Handle the input text and return the result as rendered html
        """
        if input_text is None or input_text == '':
            return ''

        classifications = self.classificator.classify(input_text)
        request = self.log_request(input_text, classifications)
        # TODO: тут надо взять хеш или ид, прокинуть его для формирования кнопок с оценкой
        return self.render(input_text, classifications)

    @staticmethod
    @db_session
    def log_request(input_text: str, classifications: List[Classification]) -> Request:
        """
            Log the request to the database
        """
        return Request(
            text=input_text,
            hash=hashlib.md5(input_text.encode()).hexdigest(),
            created_at=datetime.now(),
            updated_at=datetime.now(),
            rating=0,
            result=[c.dict() for c in classifications]
        )

    def render(self, input_text: str, classifications: List[Classification]) -> str:
        """
            Render the input text and the classifications as html text with labels
        """
        document = self.nlp(input_text)
        try:
            document.ents = [
                document.char_span(classification.start, classification.end, classification.entity) for
                classification in classifications
            ]
        except Exception as exc:
            logging.exception(exc)
        return displacy.render(document, style="ent", options=self.options)

    def run(self):
        iface = gr.Interface(
            fn=self.handle, inputs=gr.Textbox(
                lines=5, placeholder="Enter your text here",
                label='Check your text for compliance with the NVC rules'),
            outputs=["html"], examples=self.examples
        )
        iface.launch()


if __name__ == '__main__':
    db.bind(
        provider='postgres',
        user=os.getenv('pg_user'),
        password=os.getenv('pg_password'),
        host=os.getenv('pg_host'),
        port=os.getenv('pg_port'),
        database=os.getenv('pg_database')
    )
    db.generate_mapping(create_tables=True)
    application = Application(
        classificator=Classificator(
            config={
                'auth_endpoint_token': os.getenv("auth_endpoint_token"),
                'endpoint_url': os.getenv("endpoint_url")
            }
        ),
        options={"ents": ["Observation", "Evaluation"], "colors": {"Observation": "#9bddff", "Evaluation": "#f08080"}}
    )
    application.run()