|
|
import torch |
|
|
import torch.nn as nn |
|
|
from transformers import AutoTokenizer, AutoModel, T5Tokenizer, T5EncoderModel |
|
|
from transformers.modeling_outputs import BaseModelOutput |
|
|
|
|
|
try: |
|
|
import torch_npu |
|
|
from torch_npu.contrib import transfer_to_npu |
|
|
DEVICE_TYPE = "npu" |
|
|
except ModuleNotFoundError: |
|
|
DEVICE_TYPE = "cuda" |
|
|
|
|
|
|
|
|
class TransformersTextEncoderBase(nn.Module): |
|
|
""" |
|
|
Base class for text encoding using HuggingFace Transformers models. |
|
|
|
|
|
""" |
|
|
def __init__(self, model_name: str): |
|
|
super().__init__() |
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
self.model = AutoModel.from_pretrained(model_name) |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
text: list[str], |
|
|
): |
|
|
device = self.model.device |
|
|
batch = self.tokenizer( |
|
|
text, |
|
|
max_length=self.tokenizer.model_max_length, |
|
|
padding=True, |
|
|
truncation=True, |
|
|
return_tensors="pt" |
|
|
) |
|
|
input_ids = batch.input_ids.to(device) |
|
|
attention_mask = batch.attention_mask.to(device) |
|
|
output: BaseModelOutput = self.model( |
|
|
input_ids=input_ids, attention_mask=attention_mask |
|
|
) |
|
|
output = output.last_hidden_state |
|
|
mask = (attention_mask == 1).to(device) |
|
|
|
|
|
return {"output": output, "mask": mask} |
|
|
|
|
|
|
|
|
class T5TextEncoder(TransformersTextEncoderBase): |
|
|
""" |
|
|
Text encoder using T5 encoder model. |
|
|
""" |
|
|
def __init__(self, model_name: str = "/mnt/petrelfs/zhengzihao/cache/google-flan-t5-large"): |
|
|
nn.Module.__init__(self) |
|
|
self.tokenizer = T5Tokenizer.from_pretrained(model_name) |
|
|
self.model = T5EncoderModel.from_pretrained(model_name) |
|
|
for param in self.model.parameters(): |
|
|
param.requires_grad = False |
|
|
self.eval() |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
text: list[str], |
|
|
): |
|
|
with torch.no_grad(), torch.amp.autocast( |
|
|
device_type=DEVICE_TYPE, enabled=False |
|
|
): |
|
|
return super().forward(text) |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
text_encoder = T5TextEncoder() |
|
|
text = ["dog barking and cat moving"] |
|
|
text_encoder.eval() |
|
|
with torch.no_grad(): |
|
|
output = text_encoder(text) |
|
|
print(output["output"].shape) |
|
|
|
