app.py

from typing import Dict, List, Optional

import gradio

from entity_extraction.inference import NERInference
from name_matching.inference import NameMatchingInference
from utils.string_utils import Entities


def predict(narration: str, bank_name: str, company_id: str) -> Optional[List[str]]:
    prediction: List[str] = []

    entity_map: Optional[Dict[str, List[str]]] = NERInference.get_entity_map(transaction=narration, bank_name=bank_name)

    # no model found for the bank
    if entity_map is None:
        return prediction

    payer = None
    if str(Entities.PAYER.value) in entity_map:
        payer = entity_map[str(Entities.PAYER.value)][0]

    highest_score = -1

    if payer is not None:
        matches: List[str] = NameMatchingInference.get_name_matches(payer=payer, company_id=company_id)

        # no models found for the given company
        if matches is None:
            return prediction

        prediction = matches

    return prediction


gradio_interface = gradio.Interface(
    fn=predict,
    inputs=[
        gradio.Textbox(label="Narration"),
        gradio.Textbox(label="Bank Name"),
        gradio.Textbox(label="Company ID"),
    ],
    outputs=gradio.JSON()
)
gradio_interface.launch()