Valeriy Sinyukov commited on
Commit
cfad756
·
1 Parent(s): 351e3b7

Add simple pretrained model: oracat/bert-paper-classifier-arxiv

Browse files

This model predicts arXiv taxonomy: https://arxiv.org/category_taxonomy.
Link to model: https://huggingface.co/oracat/bert-paper-classifier-arxiv.

src/category_classification/models/oracat__bert_paper_classifier_arxiv/model.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+
3
+ name = "oracat/bert-paper-classifier-arxiv"
4
+
5
+
6
+ class BertPaperClassifierArxivModel:
7
+ def __init__(self):
8
+ self.pipeline = pipeline("text-classification", model=name)
9
+
10
+ def __call__(self, input):
11
+ return self.pipeline(input.title + ' ' + input.abstract)
12
+
13
+ def get_model():
14
+ return BertPaperClassifierArxivModel()