File size: 1,517 Bytes
88e3f4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import pytest
from omniff.router.encoder_router import ROUTE_DESCRIPTIONS, ROUTE_LABELS, EncoderRouter
def test_route_labels_match_descriptions():
for label in ROUTE_LABELS:
assert label in ROUTE_DESCRIPTIONS
def test_modality_overrides():
router = EncoderRouter()
assert router.route("x", input_modality="video").route_class == "VIDEO_CAPTION"
assert router.route("x", input_modality="document").route_class == "DOCUMENT_READ"
assert router.route("", input_modality="audio").route_class == "AUDIO_TRANSCRIBE_ONLY"
assert router.route("question", input_modality="audio").route_class == "AUDIO_QA"
assert (
router.route("x", input_modality="image", output_modality="image").route_class
== "IMAGE_EDIT"
)
assert router.route("x", input_modality="image").route_class == "IMAGE_CAPTION"
def test_encoder_route_text():
try:
import sentence_transformers # noqa: F401
except ImportError:
pytest.skip("sentence-transformers not installed")
router = EncoderRouter()
decision = router.route("hello how are you")
assert decision.route_class in ROUTE_LABELS
assert decision.confidence > 0
def test_encoder_route_code_detection():
try:
import sentence_transformers # noqa: F401
except ImportError:
pytest.skip("sentence-transformers not installed")
router = EncoderRouter()
decision = router.route("write a Python function that sorts a list")
assert decision.route_class == "CODE"
|