lib/model.py

import os

import tensorflow as tf
import keras
from huggingface_hub import login, hf_hub_download

from lib.result import Result

os.environ["KERAS_BACKEND"] = "jax"

login(token=os.getenv("HF_TOKEN"))


class Model:
    
    @staticmethod
    def get_english_model() -> "Model":
        return Model("elsamueldev/confia-97-english", "confia-97-english.keras")
    
    @staticmethod
    def get_spanish_model() -> "Model":
        return Model("elsamueldev/confia-97-spanish", "confia-97-spanish.keras")

    def __init__(self, repo_id: str, filename: str) -> None:
        path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            local_dir="./",
            local_dir_use_symlinks=False
        )
        if path is None:
            raise RuntimeError("Model could not be downloaded")
        
        self.__model = keras.saving.load_model(path)

    def analyze(self, text: str) -> Result:
        raw_result = self.__model.predict(tf.constant([text]))
        result = round(float(raw_result[0][0]), 3)

        return Result(percentage=result)