clhuang commited on
Commit
cdaa274
1 Parent(s): 52b6d4c

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -0
README.md CHANGED
@@ -16,6 +16,9 @@ metrics:
16
  model_path = "clhuang/albert-news-classification"
17
  model = AlbertForSequenceClassification.from_pretrained(model_path)
18
  tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
 
 
 
19
 
20
  # get category probability
21
  def get_category_proba( text ):
 
16
  model_path = "clhuang/albert-news-classification"
17
  model = AlbertForSequenceClassification.from_pretrained(model_path)
18
  tokenizer = BertTokenizer.from_pretrained("bert-base-chinese")
19
+ # Category index
20
+ news_categories=['政治','科技','運動','證卷','產經','娛樂','生活','國際','社會','文化','兩岸']
21
+ idx2cate = { i : item for i, item in enumerate(news_categories)}
22
 
23
  # get category probability
24
  def get_category_proba( text ):