File size: 1,769 Bytes
d4f8959 | 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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | 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():
# Foxconn (sentence 1) and Samsung (sentence 2) co-occur closely but in
# different sentences — must NOT be related
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
|