MCPcopy Create free account
hub / github.com/AMAP-ML/EMF / mean_pooling

Function mean_pooling

trl/examples/scripts/bco.py:92–95  ·  view source on GitHub ↗
(model_output, attention_mask)

Source from the content-addressed store, hash-verified

90 """
91
92 def mean_pooling(model_output, attention_mask):
93 token_embeddings = model_output[0]
94 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
95 return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
96
97 with torch.no_grad():
98 model_output = model(input_ids=input_ids, attention_mask=attention_mask)

Callers 1

embed_promptFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected