mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-08-09 04:57:44 +00:00
9890bb4b00
* docs: update explanation for separate config use in litellm * docs: update default backend to `rdagent.oai.backend.LiteLLMAPIBackend` * docs: update .rst format * Update installation_and_configuration.rst
33 lines
963 B
Python
33 lines
963 B
Python
import unittest
|
|
|
|
from rdagent.oai.llm_utils import (
|
|
APIBackend,
|
|
calculate_embedding_distance_between_str_list,
|
|
)
|
|
|
|
|
|
class TestEmbedding(unittest.TestCase):
|
|
def test_embedding(self) -> None:
|
|
emb = APIBackend().create_embedding("hello")
|
|
assert emb is not None
|
|
assert isinstance(emb, list)
|
|
assert len(emb) > 0
|
|
|
|
def test_embedding_list(self) -> None:
|
|
emb = APIBackend().create_embedding(["hello", "hi"])
|
|
assert emb is not None
|
|
assert isinstance(emb, list)
|
|
assert len(emb) == 2
|
|
|
|
def test_embedding_similarity(self) -> None:
|
|
similarity = calculate_embedding_distance_between_str_list(["Hello"], ["Hi"])[0][0]
|
|
assert similarity is not None
|
|
assert isinstance(similarity, float)
|
|
min_similarity_threshold = 0.8
|
|
print(f"similarity: {similarity}")
|
|
assert similarity >= min_similarity_threshold
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|