Spaces:
Sleeping
Sleeping
File size: 1,118 Bytes
2d22962 5db7813 6b5fdb9 5db7813 2d22962 d57f40c 2d22962 d57f40c 5db7813 6b5fdb9 2d22962 6b5fdb9 2d22962 d57f40c 6b5fdb9 5db7813 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel, AutoConfig
from utils import get_dataset
class Model(nn.Module):
def __init__(self, model_name, config, num_labels):
super().__init__()
self.model = AutoModel.from_pretrained(model_name, config=config)
self.classification_head = nn.Linear(config.hidden_size, num_labels)
def forward(self, input_ids):
outputs = self.model(input_ids)
pooled_output = outputs.last_hidden_state[:, 0]
logits = self.classification_head(pooled_output)
probabilities = F.softmax(logits, dim=-1)
return probabilities
def main():
model_name = "moussaKam/AraBART"
tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)
dataset = get_dataset("data/DA_train_labeled.tsv", "data/DA_dev_labeled.tsv", tokenizer)
num_labels = len(set(dataset["train"]["label"]))
model = Model(model_name, config, num_labels)
print(dataset["train"])
if __name__ == "__main__":
main() |