import torch
from torch import Tensor
from torch import nn
from typing import Union, Tuple, List, Iterable, Dict
import os
import json


class LayerNorm(nn.Module):
    def __init__(self, dimension: int):
        super(LayerNorm, self).__init__()
        self.dimension = dimension
        self.norm = nn.LayerNorm(dimension)
   

    def forward(self, features: Dict[str, Tensor]):
        features['sentence_embedding'] = self.norm(features['sentence_embedding']) 
        return features


    def get_sentence_embedding_dimension(self):
        return self.dimension

    def save(self, output_path):
        with open(os.path.join(output_path, 'config.json'), 'w') as fOut:
            json.dump({'dimension': self.dimension}, fOut, indent=2)

        torch.save(self.state_dict(), os.path.join(output_path, 'pytorch_model.bin'))

    @staticmethod
    def load(input_path):
        with open(os.path.join(input_path, 'config.json')) as fIn:
            config = json.load(fIn)

        model = LayerNorm(**config)
        model.load_state_dict(torch.load(os.path.join(input_path, 'pytorch_model.bin'), map_location=torch.device('cpu')))
        return model