Spaces:
Runtime error
Runtime error
Update app.py
Browse files
app.py
CHANGED
@@ -3,7 +3,7 @@ import pandas as pd
|
|
3 |
import torch
|
4 |
from torch import nn
|
5 |
from torch.utils.data import Dataset, DataLoader
|
6 |
-
from transformers import AutoTokenizer, BertModel
|
7 |
from sklearn import metrics
|
8 |
import streamlit as st
|
9 |
|
@@ -90,16 +90,14 @@ class BertClass(torch.nn.Module):
|
|
90 |
class PretrainedBertClass(torch.nn.Module):
|
91 |
def __init__(self):
|
92 |
super(PretrainedBertClass, self).__init__()
|
93 |
-
self.l1 =
|
94 |
-
|
95 |
-
|
96 |
-
|
97 |
-
|
98 |
-
|
99 |
-
|
100 |
-
|
101 |
-
return_dict=False,
|
102 |
-
)
|
103 |
return output
|
104 |
|
105 |
# User selects model for front-end.
|
|
|
3 |
import torch
|
4 |
from torch import nn
|
5 |
from torch.utils.data import Dataset, DataLoader
|
6 |
+
from transformers import AutoTokenizer, BertModel
|
7 |
from sklearn import metrics
|
8 |
import streamlit as st
|
9 |
|
|
|
90 |
class PretrainedBertClass(torch.nn.Module):
|
91 |
def __init__(self):
|
92 |
super(PretrainedBertClass, self).__init__()
|
93 |
+
self.l1 = BertModel.from_pretrained(bert_path)
|
94 |
+
self.l2 = torch.nn.Dropout(0.3)
|
95 |
+
self.l3 = torch.nn.Linear(768, 6)
|
96 |
+
|
97 |
+
def forward(self, ids, mask, token_type_ids):
|
98 |
+
_, output_1= self.l1(ids, attention_mask = mask, token_type_ids = token_type_ids)
|
99 |
+
output_2 = self.l2(output_1)
|
100 |
+
output = self.l3(output_2)
|
|
|
|
|
101 |
return output
|
102 |
|
103 |
# User selects model for front-end.
|