add disk embedding api (#11585)
This commit is contained in:
		
							parent
							
								
									79c742dfd5
								
							
						
					
					
						commit
						c279849d27
					
				
					 1 changed files with 19 additions and 2 deletions
				
			
		| 
						 | 
				
			
			@ -88,10 +88,11 @@ class DiskEmbedding(torch.nn.Embedding):
 | 
			
		|||
                 sparse: bool = False,
 | 
			
		||||
                 _weight: Optional[Tensor] = None,
 | 
			
		||||
                 _freeze: bool = False,
 | 
			
		||||
                 device=None, dtype=None) -> None:
 | 
			
		||||
                 device=None,
 | 
			
		||||
                 dtype=None):
 | 
			
		||||
        super().__init__(num_embeddings, embedding_dim, padding_idx,
 | 
			
		||||
                         max_norm, norm_type, scale_grad_by_freq,
 | 
			
		||||
                         sparse, _weight, _freeze, device, dtype)
 | 
			
		||||
                         sparse, _weight, True, device, dtype)
 | 
			
		||||
        self.filename = "embeddings.bin"
 | 
			
		||||
        self.weight.data.flatten().half().numpy().tofile(self.filename)
 | 
			
		||||
        dummy_weight = torch.empty(0, 0, dtype=self.weight.dtype, device=self.weight.device)
 | 
			
		||||
| 
						 | 
				
			
			@ -118,6 +119,22 @@ class DiskEmbedding(torch.nn.Embedding):
 | 
			
		|||
        )
 | 
			
		||||
        self.weight = torch.nn.Parameter(embeds, requires_grad=False)
 | 
			
		||||
 | 
			
		||||
    @classmethod
 | 
			
		||||
    def from_embedding(cls, embedding: torch.nn.Embedding):
 | 
			
		||||
        return cls(
 | 
			
		||||
            embedding.num_embeddings,
 | 
			
		||||
            embedding.embedding_dim,
 | 
			
		||||
            embedding.padding_idx,
 | 
			
		||||
            embedding.max_norm,
 | 
			
		||||
            embedding.norm_type,
 | 
			
		||||
            embedding.scale_grad_by_freq,
 | 
			
		||||
            embedding.sparse,
 | 
			
		||||
            embedding.weight.data,
 | 
			
		||||
            True,
 | 
			
		||||
            embedding.weight.device,
 | 
			
		||||
            embedding.weight.dtype,
 | 
			
		||||
        )
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class LowBitEmbedding(torch.nn.Embedding):
 | 
			
		||||
    def __init__(self,
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
		Loading…
	
		Reference in a new issue