import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

# Define mean pooling function
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

# Create a wrapper class for tracing
class TransformerWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        
    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        pooled = mean_pooling(outputs, attention_mask)
        return F.normalize(pooled, p=2, dim=1)

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('nomic-ai/nomic-embed-text-v1')
tokenizer.model_max_length = 128
base_model = AutoModel.from_pretrained('nomic-ai/nomic-embed-text-v1', trust_remote_code=True)
base_model.eval()

# Create wrapped model
wrapped_model = TransformerWrapper(base_model)

# Prepare example input for tracing
example_sentences = ['example sentence']
encoded_input = tokenizer(
    example_sentences,
    padding="max_length",
    truncation=True,
    return_tensors='pt'
)

with torch.no_grad():
    output = wrapped_model(encoded_input["input_ids"], encoded_input["attention_mask"])


# Trace the model
with torch.no_grad():
    traced_model = torch.jit.trace(
        wrapped_model,
        (
            encoded_input['input_ids'],
            encoded_input['attention_mask']
        )
    )

print(type(base_model))
      
torch.jit.save(traced_model, "my_model.pt")
torch.jit.load("my_model.pt") # this will fail

$ print(type(base_model))
<class 'transformers_modules.nomic-ai.nomic-bert-2048.40b98394640e630d5276807046089b233113aa87.modeling_hf_nomic_bert.NomicBertModel'>`

    return torch.jit.load(model_path)
    cpp_module = torch._C.import_ir_module(cu, str(f), map_location, _extra_files, _restore_shapes)  # type: ignore[call-arg]
RuntimeError: expected newline but found 'ident' here:
Serialized   File "code/__torch__.py", line 6
  training : bool
  _is_full_backward_hook : Optional[bool]
  model : __torch__.transformers_modules.nomic-ai.nomic-bert-2048.40b98394640e630d5276807046089b233113aa87.modeling_hf_nomic_bert.NomicBertModel
                                                                    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  def forward(self: __torch__.TransformerWrapper,
    input_ids: Tensor,

print(type(base_model))
<class 'transformers_modules.nomic_ai.nomic_bert_2048._40b98394640e630d5276807046089b233113aa87.modeling_hf_nomic_bert.NomicBertModel'>