Eric22333 commited on
Commit
8f4a6b9
·
verified ·
1 Parent(s): 9b76705

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +3 -3
pipeline.py CHANGED
@@ -6,7 +6,7 @@ import numpy as np
6
  import json
7
 
8
  class NewsClassifierConfig(PretrainedConfig):
9
- model_type = "news_classifier"
10
 
11
  def __init__(
12
  self,
@@ -26,7 +26,7 @@ class NewsClassifierConfig(PretrainedConfig):
26
 
27
  class NewsClassifier(PreTrainedModel):
28
  config_class = NewsClassifierConfig
29
- base_model_prefix = "news_classifier"
30
 
31
  def __init__(self, config):
32
  super().__init__(config)
@@ -35,7 +35,7 @@ class NewsClassifier(PreTrainedModel):
35
 
36
  def post_init(self):
37
  """Load model and tokenizer after initialization"""
38
- self.model = load_model('news_classifier.h5')
39
  with open('tokenizer.json', 'r') as f:
40
  tokenizer_data = json.load(f)
41
  self.tokenizer = tokenizer_from_json(tokenizer_data)
 
6
  import json
7
 
8
  class NewsClassifierConfig(PretrainedConfig):
9
+ model_type = "custom"
10
 
11
  def __init__(
12
  self,
 
26
 
27
  class NewsClassifier(PreTrainedModel):
28
  config_class = NewsClassifierConfig
29
+ base_model_prefix = "custom"
30
 
31
  def __init__(self, config):
32
  super().__init__(config)
 
35
 
36
  def post_init(self):
37
  """Load model and tokenizer after initialization"""
38
+ self.model = load_model('custom.h5')
39
  with open('tokenizer.json', 'r') as f:
40
  tokenizer_data = json.load(f)
41
  self.tokenizer = tokenizer_from_json(tokenizer_data)