Skip to content

Commit 2beef1a

Browse files
committed
Compute embedding locally to avoid discrepency due to device
1 parent 1fac541 commit 2beef1a

File tree

3 files changed

+25
-393
lines changed

3 files changed

+25
-393
lines changed

src/giskard_lmutils/model/litellm.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ def get_embedding(self, input: str):
5757

5858
with torch.no_grad():
5959
outputs = self.model(**inputs)
60+
6061
return outputs.last_hidden_state.mean(dim=1).squeeze(0)
6162

6263

tests/data/embedding.py

Lines changed: 0 additions & 388 deletions
This file was deleted.

0 commit comments

Comments
 (0)