models/storm_magic/model.py

import torch, os
from safetensors import safe_open


class EmbeddingModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        model_path = os.path.join(os.path.dirname(__file__), TEMPLATE_MODEL_PATH)
        with safe_open(model_path, framework="pt", device="cpu") as f:
            shape = f.get_tensor("weight").shape
        self.weight = torch.nn.Parameter(torch.empty(shape))
        
    @torch.no_grad()
    def process_inputs(self, **kwargs):
        return {}

    def forward(self, **kwargs):
        return {"text_embedding": self.weight}


class DataProcessor:
    def __call__(self, **kwargs):
        return kwargs


TEMPLATE_MODEL = EmbeddingModel
TEMPLATE_MODEL_PATH = "model.safetensors"
TEMPLATE_DATA_PROCESSOR = DataProcessor