AlanRobotics commited on
Commit
6bff269
·
1 Parent(s): 7b8d8df

Upload model

Browse files
Files changed (2) hide show
  1. configuration_siamese.py +1 -1
  2. modeling_siamese.py +2 -0
configuration_siamese.py CHANGED
@@ -7,7 +7,7 @@ class SiameseConfig(PretrainedConfig):
7
  def __init__(
8
  self,
9
  **kwargs):
10
- super().__init__()
11
 
12
 
13
  siamese_config = SiameseConfig()
 
7
  def __init__(
8
  self,
9
  **kwargs):
10
+ super().__init__(**kwargs)
11
 
12
 
13
  siamese_config = SiameseConfig()
modeling_siamese.py CHANGED
@@ -1,5 +1,6 @@
1
  from transformers import PreTrainedModel, BertModel
2
  import torch
 
3
 
4
  checkpoint = 'cointegrated/rubert-tiny'
5
 
@@ -36,6 +37,7 @@ second_model = SiameseNN()
36
  second_model.load_state_dict(torch.load('siamese_state'))
37
 
38
  class SiamseNNModel(PreTrainedModel):
 
39
  def __init__(self, config):
40
  super().__init__(config)
41
  self.model = second_model
 
1
  from transformers import PreTrainedModel, BertModel
2
  import torch
3
+ from .configuration_siamese import SiameseConfig
4
 
5
  checkpoint = 'cointegrated/rubert-tiny'
6
 
 
37
  second_model.load_state_dict(torch.load('siamese_state'))
38
 
39
  class SiamseNNModel(PreTrainedModel):
40
+ config_class = SiameseConfig
41
  def __init__(self, config):
42
  super().__init__(config)
43
  self.model = second_model