| from backend.relation_extractor import extract_patterns, extract_relations, _find_mentions |
|
|
|
|
| CHUNK = ( |
| "Foxconn is a major supplier to Apple Inc. for iPhone assembly. " |
| "In contrast, Apple competes with Samsung in the smartphone market. " |
| "Beats Electronics, a subsidiary of Apple, also contributed to revenue." |
| ) |
| ORGS = ["Foxconn", "Apple Inc.", "Samsung", "Beats Electronics", "Apple"] |
|
|
|
|
| def _rels(pairs): |
| return {(a.lower(), b.lower()): r for a, b, r in pairs} |
|
|
|
|
| def test_supplier_pattern_found(): |
| rels = _rels(extract_patterns(CHUNK, ORGS)) |
| assert rels.get(("foxconn", "apple inc.")) == "SUPPLIER_TO" |
|
|
|
|
| def test_competitor_pattern_found(): |
| rels = _rels(extract_patterns(CHUNK, ORGS)) |
| assert rels.get(("apple", "samsung")) == "COMPETITOR_OF" |
|
|
|
|
| def test_subsidiary_pattern_found(): |
| rels = _rels(extract_patterns(CHUNK, ORGS)) |
| assert rels.get(("beats electronics", "apple")) == "SUBSIDIARY_OF" |
|
|
|
|
| def test_cross_sentence_pairs_not_connected(): |
| |
| |
| rels = _rels(extract_patterns(CHUNK, ORGS)) |
| assert ("foxconn", "samsung") not in rels |
| assert ("samsung", "foxconn") not in rels |
|
|
|
|
| def test_overlapping_mentions_collapse_to_longest(): |
| spans = _find_mentions("Apple Inc. was mentioned.", ["Apple", "Apple Inc."]) |
| assert len(spans) == 1 |
| assert spans[0][2] == "Apple Inc." |
|
|
|
|
| def test_single_org_returns_empty(): |
| assert extract_relations(CHUNK, ["Apple"], use_llm_fallback=False) == [] |
|
|
|
|
| def test_pattern_only_mode_needs_no_llm(): |
| results = extract_relations(CHUNK, ORGS, use_llm_fallback=False) |
| assert all(r["source"] == "pattern" for r in results) |
| assert len(results) >= 3 |
|
|
