|
|
import json |
|
|
import os |
|
|
from typing import Union, List, Dict, Tuple |
|
|
|
|
|
import torch |
|
|
from sentence_transformers import models |
|
|
from transformers import AutoModel |
|
|
|
|
|
|
|
|
class EmbeddingModel(models.Transformer): |
|
|
def __init__(self, *args, **kwargs): |
|
|
self.model_name_or_path = "lamarr-llm-development/elbedding" |
|
|
kwargs.pop("model_name_or_path", None) |
|
|
super().__init__(*args, **kwargs) |
|
|
|
|
|
def tokenize( |
|
|
self, |
|
|
texts: Union[List[str], List[Dict], List[Tuple[str, str]]], |
|
|
padding: Union[str, bool] = True, |
|
|
) -> Dict[str, torch.Tensor]: |
|
|
"""Tokenizes a text and maps tokens to token-ids""" |
|
|
output = {} |
|
|
if isinstance(texts[0], str): |
|
|
texts = [x + self.tokenizer.eos_token for x in texts] |
|
|
to_tokenize = [texts] |
|
|
elif isinstance(texts[0], dict): |
|
|
to_tokenize = [] |
|
|
output["text_keys"] = [] |
|
|
for lookup in texts: |
|
|
text_key, text = next(iter(lookup.items())) |
|
|
to_tokenize.append(text) |
|
|
output["text_keys"].append(text_key) |
|
|
to_tokenize = [to_tokenize] |
|
|
else: |
|
|
batch1, batch2 = [], [] |
|
|
for text_tuple in texts: |
|
|
batch1.append(text_tuple[0]) |
|
|
batch2.append(text_tuple[1]) |
|
|
to_tokenize = [batch1, batch2] |
|
|
|
|
|
output.update( |
|
|
self.tokenizer( |
|
|
*to_tokenize, |
|
|
padding="max_length", |
|
|
truncation=True, |
|
|
return_tensors="pt", |
|
|
max_length=512, |
|
|
) |
|
|
) |
|
|
|
|
|
|
|
|
output.pop("token_type_ids", None) |
|
|
|
|
|
return output |
|
|
|
|
|
def get_config_dict(self) -> dict[str, str]: |
|
|
return {"model_name_or_path": self.model_name_or_path} |
|
|
|
|
|
def save(self, save_dir: str, **kwargs) -> None: |
|
|
self.auto_model.save_pretrained(save_dir, safe_serialization=True) |
|
|
self.tokenizer.save_pretrained(save_dir) |
|
|
|
|
|
with open(os.path.join(save_dir, "sentence_bert_config.json"), "w+") as f: |
|
|
json.dump(self.get_config_dict(), f, indent=4) |
|
|
|
|
|
@staticmethod |
|
|
def load(**kwargs) -> "EmbeddingModel": |
|
|
return EmbeddingModel(**kwargs) |
|
|
|
