| import torch.nn as nn | |
| from transformers import AutoModel | |
| class ThaiEncoder(nn.Module): | |
| def __init__(self, model_name: str, trainable: bool = False) -> None: | |
| super().__init__() | |
| self.model = AutoModel.from_pretrained(model_name) | |
| for p in self.model.parameters(): | |
| p.requires_grad = trainable | |
| self.target_token_idx = 0 | |
| def forward(self, input_ids, attention_mask): | |
| output = self.model(input_ids=input_ids, attention_mask=attention_mask) | |
| last_hidden_state = output.last_hidden_state | |
| return last_hidden_state[:, self.target_token_idx, :] |