File size: 770 Bytes
99f8f6e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | 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
|