Always convert YAML tags to strings

This commit is contained in:
Augusto de la Torre 2023-04-23 00:36:32 +02:00
parent 9b6c01a7f0
commit 2e8056174d
3 changed files with 4 additions and 3 deletions

View File

@ -36,7 +36,7 @@ class Tag:
return Tag(data)
if isinstance(data, dict):
value = data.get("tag")
value = str(data.get("tag"))
weight = data.get("weight")
if value:
return Tag(value, weight)

View File

@ -131,7 +131,7 @@ class TestResolve(unittest.TestCase):
self.assertEqual(len(items), 6)
self.assertEqual(set(image_paths), set([IMAGE_1_PATH, IMAGE_2_PATH, IMAGE_3_PATH] * 2))
self.assertEqual(set(captions), {}'caption for test1', 'test2', 'test3', 'caption for test1', 'caption for test2', 'test3'})
self.assertEqual(set(captions), {'caption for test1', 'test2', 'test3', 'caption for test1', 'caption for test2', 'test3'})
undersized_images = list(filter(lambda i: i.is_undersized, items))
self.assertEqual(len(undersized_images), 2)

View File

@ -77,6 +77,7 @@ class TestDataset(TestCase):
- tag: from .yml
- tag: with weight
weight: 0.5
- tag: 1234.5
"""))
actual = Dataset.from_path(".").image_configs
@ -94,7 +95,7 @@ class TestDataset(TestCase):
rating=1.1,
max_caption_length=1024,
main_prompts="A complex caption",
tags= { Tag("from .yml"), Tag("with weight", weight=0.5) }
tags= { Tag("from .yml"), Tag("with weight", weight=0.5), Tag("1234.5") }
)
}
self.assertEqual(expected, actual)