DummyModel.py

import torch
import torch.nn


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)