Files
Auto-GPT/tests/embedder_test.py
Tymec 062176d3f5 test: replaced MockConfig with real config
get_embedding function uses config
2023-04-14 22:49:13 +02:00

30 lines
685 B
Python

import os
import sys
# Probably a better way:
sys.path.append(os.path.abspath('../scripts'))
from memory.base import get_embedding
from config import Config
import unittest
# Required, because the get_embedding function uses it
cfg = Config()
class TestMemoryEmbedder(unittest.TestCase):
def test_ada(self):
cfg.memory_embedder = "ada"
text = "Sample text"
result = get_embedding(text)
self.assertEqual(len(result), 1536)
def test_sbert(self):
cfg.memory_embedder = "sbert"
text = "Sample text"
result = get_embedding(text)
self.assertEqual(len(result), 768)
if __name__ == '__main__':
unittest.main()