Valeriy Sinyukov commited on
Commit
4d7b773
·
1 Parent(s): cfad756

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

Browse files

This model predicts category for biomedical papers.
Link to model: https://huggingface.co/oracat/bert-paper-classifier.

src/category_classification/models/oracat__bert_paper_classifier/model.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+
3
+ name = "oracat/bert-paper-classifier"
4
+
5
+
6
+ class BertPaperClassifierModel:
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 __getstate__(self):
14
+ return self.__dict__
15
+
16
+ def get_model():
17
+ return BertPaperClassifierModel()