DummyModel.py

import torch
import os
import torch.nn
from huggingface_hub import hf_hub_download


def load_dummy_model(DEBUG):
    model = DummyModel()
    if not DEBUG:
        file_path = hf_hub_download(
            "lfolle/DeepNAPSIModel",
            "dummy_model.pth",
            use_auth_token=os.environ["DeepNAPSIModel"],
        )
        model.load_state_dict(torch.load(file_path))
    return model


class DummyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: list):
        return torch.softmax(torch.rand(len(x), 5), 1), 0

    def __call__(self, x: list):
        return self.forward(x)