| # model/luna_model.py | |
| import torch | |
| import torch.nn as nn | |
| from transformers import BertTokenizer, BertModel | |
| class LunaAI(nn.Module): | |
| def __init__(self): | |
| super(LunaAI, self).__init__() | |
| self.bert = BertModel.from_pretrained('bert-base-uncased') | |
| self.classifier = nn.Linear(768, 2) # Adjust for number of classes | |
| def forward(self, input_ids, attention_mask): | |
| outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) | |
| logits = self.classifier(outputs.pooler_output) | |
| return logits | |