mirror of
https://github.com/MrSleeps/jellyfin-scene-search.git
synced 2026-11-14 22:11:27 +00:00
40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from scene_search.embedder import SiglipEmbedder
|
|
|
|
|
|
class SiglipMatchScoringTests(unittest.TestCase):
|
|
def test_match_probabilities_use_model_scale_and_bias(self) -> None:
|
|
embedder = SiglipEmbedder.__new__(SiglipEmbedder)
|
|
embedder.device = torch.device("cpu")
|
|
embedder.model = SimpleNamespace(
|
|
logit_scale=torch.tensor(math.log(2.0)),
|
|
logit_bias=torch.tensor(-1.0),
|
|
)
|
|
|
|
text = np.asarray([1.0, 0.0], dtype=np.float32)
|
|
images = np.asarray(
|
|
[
|
|
[1.0, 0.0],
|
|
[0.0, 1.0],
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
|
|
scores = embedder.match_probabilities(text, images)
|
|
|
|
self.assertAlmostEqual(float(scores[0]), 0.7310586, places=6)
|
|
self.assertAlmostEqual(float(scores[1]), 0.2689414, places=6)
|
|
self.assertGreater(float(scores[0]), float(scores[1]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|