Text Generation
Transformers
PyTorch
mpt
Composer
MosaicML
llm-foundry
custom_code
text-generation-inference
mpt-7b-instruct / custom_embedding.py
daking's picture
LLM-foundry update June 16, 2023 22:55:57
4f5f899
raw
history blame
305 Bytes
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
class SharedEmbedding(nn.Embedding):
def forward(self, input: Tensor, unembed: bool=False) -> Tensor:
if unembed:
return F.linear(input, self.weight)
return super().forward(input)