Update app.py

app.py

@@ -17,7 +17,8 @@ idx_to_tag = {0: 'cs',
  3: 'math',
  4: 'q-bio',
  5: 'eess',
- 6: 'economics, finances'}
+ 6: 'economics, finances'
+             }
 
 
 tag_to_idx = {'cs': 0,
@@ -29,24 +30,6 @@ tag_to_idx = {'cs': 0,
  'economics, finances': 6
 }
 
-class RobertaClass(torch.nn.Module):
-    def __init__(self):
-        super(RobertaClass, self).__init__()
-        self.l1 = RobertaModel.from_pretrained("roberta-base")
-        self.pre_classifier = torch.nn.Linear(768, 768)
-        self.dropout = torch.nn.Dropout(0.3)
-        self.classifier = torch.nn.Linear(768, 10)
-
-    def forward(self, input_ids, attention_mask, token_type_ids):
-        output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
-        hidden_state = output_1[0]
-        pooler = hidden_state[:, 0]
-        pooler = self.pre_classifier(pooler)
-        pooler = torch.nn.ReLU()(pooler)
-        pooler = self.dropout(pooler)
-        output = self.classifier(pooler)
-        return output
-
 tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
 
 model = AutoModel.from_pretrained('./model')