mirror of
https://github.com/wassname/Castor.git
synced 2026-08-20 12:00:37 +08:00
30 lines
938 B
Python
Executable File
30 lines
938 B
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from han.sent_level_rnn import SentLevelRNN
|
|
from han.word_level_rnn import WordLevelRNN
|
|
|
|
|
|
class HAN(nn.Module):
|
|
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
dataset = config.dataset
|
|
self.mode = config.mode
|
|
self.word_attention_rnn = WordLevelRNN(config)
|
|
self.sentence_attention_rnn = SentLevelRNN(config)
|
|
|
|
def forward(self, x, **kwargs):
|
|
x = x.permute(1, 2, 0) # Expected : # sentences, # words, batch size
|
|
num_sentences = x.size(0)
|
|
word_attentions = None
|
|
for i in range(num_sentences):
|
|
word_attn = self.word_attention_rnn(x[i, :, :])
|
|
if word_attentions is None:
|
|
word_attentions = word_attn
|
|
else:
|
|
word_attentions = torch.cat((word_attentions, word_attn), 0)
|
|
return self.sentence_attention_rnn(word_attentions)
|
|
|