go00od commited on
Commit
35113c3
·
verified ·
1 Parent(s): b327588

Upload model_def.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. model_def.py +16 -0
model_def.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch.nn as nn
2
+
3
+
4
+ class TextClassifier(nn.Module):
5
+ def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim):
6
+ super(TextClassifier, self).__init__()
7
+ self.embedding = nn.Embedding(vocab_size, embedding_dim)
8
+ self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
9
+ self.fc = nn.Linear(hidden_dim, output_dim)
10
+
11
+ def forward(self, x):
12
+ embedded = self.embedding(x)
13
+ lstm_out, (hidden, cell) = self.lstm(embedded)
14
+ last_hidden = hidden.squeeze(0)
15
+ logits = self.fc(last_hidden)
16
+ return logits