Added auto_map to the config
Browse files- config.json +5 -1
 - modeling.py +5 -0
 
    	
        config.json
    CHANGED
    
    | 
         @@ -1,7 +1,11 @@ 
     | 
|
| 1 | 
         
             
            {
         
     | 
| 2 | 
         
             
              "architectures": [
         
     | 
| 3 | 
         
            -
                " 
     | 
| 4 | 
         
             
              ],
         
     | 
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 5 | 
         
             
              "entity_model_cfg": {
         
     | 
| 6 | 
         
             
                "aggregate_func": "sum",
         
     | 
| 7 | 
         
             
                "hidden_dims": [
         
     | 
| 
         | 
|
| 1 | 
         
             
            {
         
     | 
| 2 | 
         
             
              "architectures": [
         
     | 
| 3 | 
         
            +
                "UltraForKnowledgeGraphReasoning"
         
     | 
| 4 | 
         
             
              ],
         
     | 
| 5 | 
         
            +
              "auto_map": {
         
     | 
| 6 | 
         
            +
                "AutoConfig": "modeling.UltraConfig",
         
     | 
| 7 | 
         
            +
                "AutoModel": "modeling.UltraForKnowledgeGraphReasoning"
         
     | 
| 8 | 
         
            +
              },
         
     | 
| 9 | 
         
             
              "entity_model_cfg": {
         
     | 
| 10 | 
         
             
                "aggregate_func": "sum",
         
     | 
| 11 | 
         
             
                "hidden_dims": [
         
     | 
    	
        modeling.py
    CHANGED
    
    | 
         @@ -11,6 +11,11 @@ class UltraConfig(PretrainedConfig): 
     | 
|
| 11 | 
         | 
| 12 | 
         
             
                model_type = "ultra"
         
     | 
| 13 | 
         | 
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 
         | 
|
| 14 | 
         
             
                def __init__(
         
     | 
| 15 | 
         
             
                        self,
         
     | 
| 16 | 
         
             
                        relation_model_layers: int = 6,
         
     | 
| 
         | 
|
| 11 | 
         | 
| 12 | 
         
             
                model_type = "ultra"
         
     | 
| 13 | 
         | 
| 14 | 
         
            +
                auto_map = {
         
     | 
| 15 | 
         
            +
                    "AutoConfig": "modeling.UltraConfig",
         
     | 
| 16 | 
         
            +
                    "AutoModel": "modeling.UltraForKnowledgeGraphReasoning",
         
     | 
| 17 | 
         
            +
                }
         
     | 
| 18 | 
         
            +
             
     | 
| 19 | 
         
             
                def __init__(
         
     | 
| 20 | 
         
             
                        self,
         
     | 
| 21 | 
         
             
                        relation_model_layers: int = 6,
         
     |