Update app.py
Browse files
app.py
CHANGED
@@ -15,7 +15,7 @@ tokenizer = get_tokenizer()
|
|
15 |
class devops_model(nn.Module):
|
16 |
def __init__(self):
|
17 |
super(devops_model, self).__init__()
|
18 |
-
self.
|
19 |
self.fc = nn.Sequential(
|
20 |
nn.Linear(768, 768),
|
21 |
nn.ReLU(),
|
@@ -26,7 +26,7 @@ class devops_model(nn.Module):
|
|
26 |
)
|
27 |
|
28 |
def forward(self, train_batch):
|
29 |
-
emb = self.
|
30 |
return self.fc(emb)
|
31 |
|
32 |
@st.cache
|
|
|
15 |
class devops_model(nn.Module):
|
16 |
def __init__(self):
|
17 |
super(devops_model, self).__init__()
|
18 |
+
self.berta = None
|
19 |
self.fc = nn.Sequential(
|
20 |
nn.Linear(768, 768),
|
21 |
nn.ReLU(),
|
|
|
26 |
)
|
27 |
|
28 |
def forward(self, train_batch):
|
29 |
+
emb = self.berta(**train_batch)['last_hidden_state'].mean(axis=1)
|
30 |
return self.fc(emb)
|
31 |
|
32 |
@st.cache
|