mirror of
https://github.com/aljazceru/Auto-GPT.git
synced 2026-01-08 16:54:31 +01:00
27 lines
613 B
Python
27 lines
613 B
Python
import os
|
|
import sys
|
|
|
|
from autogpt.config import Config
|
|
from autogpt.memory.base import get_embedding
|
|
|
|
# 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()
|