File size: 698 Bytes
5bd8e7f |
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 |
import unittest
from mlflow.parser import convert_to_knowledge_graph_spec
class TestParser(unittest.TestCase):
def test_parser(self):
model_results: list = [
[
[['我', 'ARG0', 0, 1], ['爱', 'PRED', 1, 2], ['中国', 'ARG1', 2, 3]]
]
]
expected_nodes = [
'我',
'中国'
]
expected_links = ['爱']
assert [node["fields"]["name"] for node in
convert_to_knowledge_graph_spec(model_results)["nodes"]] == expected_nodes
assert [node["fields"]["type"] for node in
convert_to_knowledge_graph_spec(model_results)["links"]] == expected_links
|