| from transformers import PretrainedConfig | |
| class SimilarityModelConfig(PretrainedConfig): | |
| model_type = 'roberta' | |
| def __init__(self, **kwargs): | |
| super().__init__(**kwargs) | |
| self.embedding_model_config = kwargs.get("embedding_model_config") | |
| self.score_model_config = kwargs.get("score_model_config") | |
| self.weighting_function_config = kwargs.get("weighting_function_config") | |
| nama_base = SimilarityModelConfig( | |
| embedding_model_config={ | |
| "model_class": 'roberta', | |
| "model_name":'roberta-base', | |
| "pooling": 'pooler', | |
| "normalize":True, | |
| "d":128, | |
| "prompt":'', | |
| "device":'cpu', | |
| "add_upper": True, | |
| "upper_case":False | |
| }, | |
| score_model_config={"alpha": 50}, | |
| weighting_function_config={"weighting_exponent": 0.5}, | |
| device="cpu", | |
| ) | |