Robin
Claude Sonnet 4.6
feat: smarter fallback โ heuristic sufficiency + result merging (v3.1)
2288fd7 Download tests/test_extract.py from RobinWu/nerserver: direct link, hf CLI and curl.
- Browser
- Download file 23.3 kB
-
https://huggingface.co/spaces/RobinWu/nerserver/resolve/main/tests/test_extract.py
- Command line
-
hf download hf://spaces/RobinWu/nerserver/tests/test_extract.py
-
curl -L -o test_extract.py https://huggingface.co/spaces/RobinWu/nerserver/resolve/main/tests/test_extract.py
23.3 kB
| """ | |
| Unit tests โ no real model loaded (GLiNER/torch stubbed via conftest.py). | |
| Covers: | |
| - API contract (health, validation, threshold, language field) | |
| - Dual-model routing (ZH โ BERT backend, EN/AR/mixed โ GLiNER backend) | |
| - Optional labels + bilingual auto-expansion | |
| - English / Chinese / Arabic / mixed-language scenarios | |
| - labels_used echo in response | |
| """ | |
| from unittest.mock import MagicMock, patch, PropertyMock | |
| import pytest | |
| from fastapi.testclient import TestClient | |
| from app.main import app | |
| from app.models import Entity | |
| from app.labels import DEFAULT_LABELS, expand_bilingual, labels_to_bert_types | |
| from app.ner import detect_language | |
| # โโ Helpers โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def _ents(*args: Entity) -> tuple[list[Entity], list[str]]: | |
| """Wrap entities in (entities, labels_used) tuple expected by NERService.""" | |
| return list(args), [e.label for e in args] | |
| # โโ Fixture โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def client(): | |
| """ | |
| Patch NERService so no model is actually loaded. | |
| mock_ner.extract() returns ([], []) by default. | |
| """ | |
| mock_ner = MagicMock() | |
| mock_ner.extract.return_value = ([], []) | |
| with pytest.MonkeyPatch().context() as mp: | |
| mp.setattr("app.main.NERService", lambda *_: mock_ner) | |
| with TestClient(app) as c: | |
| yield c, mock_ner | |
| # โโ System / API contract โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_health(client): | |
| c, _ = client | |
| assert c.get("/api/v1/health").json() == {"status": "ok"} | |
| def test_extract_empty_text_returns_empty(client): | |
| c, _ = client | |
| resp = c.post("/api/v1/extract", json={"text": "", "labels": ["person"]}) | |
| assert resp.status_code == 200 | |
| assert resp.json()["entities"] == [] | |
| def test_extract_labels_optional(client): | |
| """labels ๅญๆฎตๅฎๅ จไธไผ ๅบๆญฃๅธธ่ฟๅ 200ใ""" | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = ([], DEFAULT_LABELS) | |
| resp = c.post("/api/v1/extract", json={"text": "Some text."}) | |
| assert resp.status_code == 200 | |
| assert len(resp.json()["labels_used"]) > 0 | |
| def test_extract_empty_labels_uses_defaults(client): | |
| """labels=[] ๆถๅบไฝฟ็จ้ป่ฎคๅ่ฏญๆ ็ญพ้ใ""" | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = ([], DEFAULT_LABELS) | |
| resp = c.post("/api/v1/extract", json={"text": "Hello world.", "labels": []}) | |
| assert resp.status_code == 200 | |
| assert resp.json()["labels_used"] == DEFAULT_LABELS | |
| def test_extract_threshold_forwarded(client): | |
| c, mock_ner = client | |
| c.post("/api/v1/extract", | |
| json={"text": "Hello", "labels": ["person"], "threshold": 0.8}) | |
| mock_ner.extract.assert_called_once_with( | |
| "Hello", ["person"], 0.8, language="auto", min_entities=None | |
| ) | |
| def test_extract_invalid_threshold(client): | |
| c, _ = client | |
| assert c.post("/api/v1/extract", | |
| json={"text": "x", "threshold": 1.5}).status_code == 422 | |
| def test_extract_language_field_forwarded(client): | |
| c, mock_ner = client | |
| c.post("/api/v1/extract", | |
| json={"text": "ๅไบฌๅๅๅป้ข", "labels": ["ๅป้ขๅ็งฐ"], "language": "zh"}) | |
| mock_ner.extract.assert_called_once_with( | |
| "ๅไบฌๅๅๅป้ข", ["ๅป้ขๅ็งฐ"], 0.4, language="zh", min_entities=None | |
| ) | |
| def test_extract_min_entities_forwarded(client): | |
| c, mock_ner = client | |
| c.post("/api/v1/extract", | |
| json={"text": "้ฉฌไบๅจๆญๅทใ", "language": "zh", "min_entities": 5}) | |
| mock_ner.extract.assert_called_once_with( | |
| "้ฉฌไบๅจๆญๅทใ", [], 0.4, language="zh", min_entities=5 | |
| ) | |
| def test_extract_negative_min_entities_rejected(client): | |
| c, _ = client | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "x", "min_entities": -1}) | |
| assert resp.status_code == 422 | |
| def test_extract_invalid_language(client): | |
| c, _ = client | |
| assert c.post("/api/v1/extract", | |
| json={"text": "x", "language": "jp"}).status_code == 422 | |
| def test_labels_used_echoed(client): | |
| c, mock_ner = client | |
| used = ["ไบบๅๆๅงๅ", "ๅฐๅๆๅๅธ"] | |
| mock_ner.extract.return_value = ([], used) | |
| resp = c.post("/api/v1/extract", json={"text": "้ฉฌไบๅจๆญๅทใ", "labels": ["ไบบๅๆๅงๅ"]}) | |
| assert resp.json()["labels_used"] == used | |
| def test_entity_response_fields(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="Apple", label="organization", score=0.95, start=0, end=5) | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "Apple is great.", "labels": ["organization"]}) | |
| e = resp.json()["entities"][0] | |
| assert {"text", "label", "score", "start", "end"} <= e.keys() | |
| assert 0.0 <= e["score"] <= 1.0 | |
| assert e["start"] < e["end"] | |
| # โโ Language detection โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_detect_language_english(): | |
| assert detect_language("Elon Musk founded SpaceX in California.") == "en" | |
| def test_detect_language_chinese(): | |
| assert detect_language("้ฟ้ๅทดๅทด้ๅขๅๅงไบบ้ฉฌไบไบๆญๅทๅธไปปใ") == "zh" | |
| def test_detect_language_arabic(): | |
| assert detect_language("ุฃุนูู ุงูุฑุฆูุณ ู ุญู ุฏ ุจู ุณูู ุงู ุนู ู ุดุฑูุน ูููู .") == "ar" | |
| def test_detect_language_mixed(): | |
| assert detect_language("ๅผ ไผๅ ๅ ฅไบ Google ๅไบฌ็ ๅไธญๅฟ๏ผ่ด่ดฃ Android ไผๅใ") == "mixed" | |
| def test_detect_language_empty(): | |
| assert detect_language("") == "en" | |
| # โโ Bilingual label expansion โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_expand_bilingual_adds_english_for_chinese(): | |
| result = expand_bilingual(["ไบบๅๆๅงๅ"]) | |
| assert "full name of a person" in result | |
| def test_expand_bilingual_adds_chinese_for_english(): | |
| result = expand_bilingual(["company or organization name"]) | |
| assert "ๅ ฌๅธๆ็ป็ปๆบๆๅ็งฐ" in result | |
| def test_expand_bilingual_no_duplicate(): | |
| result = expand_bilingual(["ไบบๅๆๅงๅ", "full name of a person"]) | |
| assert result.count("ไบบๅๆๅงๅ") == 1 | |
| assert result.count("full name of a person") == 1 | |
| def test_expand_bilingual_custom_label_preserved(): | |
| result = expand_bilingual(["my custom label"]) | |
| assert "my custom label" in result | |
| def test_default_labels_bilingual(): | |
| has_en = any(all(ord(c) < 128 for c in lbl) for lbl in DEFAULT_LABELS) | |
| has_zh = any(any('ไธ' <= c <= '้ฟฟ' for c in lbl) for lbl in DEFAULT_LABELS) | |
| assert has_en and has_zh | |
| # โโ BERT label mapping โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_labels_to_bert_types_chinese_label(): | |
| types = labels_to_bert_types(["ไบบๅๆๅงๅ"]) | |
| assert "PER" in types | |
| def test_labels_to_bert_types_english_label(): | |
| types = labels_to_bert_types(["geographical location"]) | |
| assert "LOC" in types | |
| def test_labels_to_bert_types_empty_returns_none(): | |
| assert labels_to_bert_types([]) is None | |
| def test_labels_to_bert_types_unmapped_returns_none(): | |
| # ๆ ๆณๆ ๅฐ็ๆ ็ญพ โ ไธ่ฟๆปค๏ผ่ฟๅ None๏ผ | |
| assert labels_to_bert_types(["some unknown label"]) is None | |
| # โโ English โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_english_person_org(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="Elon Musk", label="person", score=0.98, start=0, end=9), | |
| Entity(text="Tesla", label="organization", score=0.96, start=18, end=23), | |
| Entity(text="SpaceX", label="organization", score=0.97, start=28, end=34), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "Elon Musk is the CEO of Tesla and founded SpaceX.", | |
| "labels": ["full name of a person", "company or organization name"], | |
| "language": "en"}) | |
| assert resp.status_code == 200 | |
| texts = {e["text"] for e in resp.json()["entities"]} | |
| assert {"Elon Musk", "Tesla", "SpaceX"} <= texts | |
| def test_english_location_date(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="Paris", label="location", score=0.94, start=20, end=25), | |
| Entity(text="2024", label="date", score=0.91, start=29, end=33), | |
| Entity(text="France", label="location", score=0.93, start=38, end=44), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "The summit was held in Paris in 2024, in France.", | |
| "labels": ["geographical location", "date or year"], | |
| "language": "en"}) | |
| texts = {e["text"] for e in resp.json()["entities"]} | |
| assert {"Paris", "France", "2024"} <= texts | |
| def test_english_threshold_forwarded(client): | |
| c, mock_ner = client | |
| c.post("/api/v1/extract", | |
| json={"text": "NASA explored the Moon.", | |
| "labels": ["company or organization name"], | |
| "threshold": 0.8, "language": "en"}) | |
| mock_ner.extract.assert_called_once_with( | |
| "NASA explored the Moon.", ["company or organization name"], 0.8, | |
| language="en", min_entities=None, | |
| ) | |
| # โโ Chinese (BERT backend) โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_chinese_person_org(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="้ฉฌไบ", label="ไบบๅๆๅงๅ", score=0.96, start=8, end=10), | |
| Entity(text="ๅผ ๅ", label="ไบบๅๆๅงๅ", score=0.94, start=25, end=27), | |
| Entity(text="้ฟ้ๅทดๅทด", label="ๅ ฌๅธๆ็ป็ปๆบๆๅ็งฐ", score=0.97, start=0, end=4), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "้ฟ้ๅทดๅทด้ๅขๅๅงไบบ้ฉฌไบๅธไปป๏ผ็ฑๅผ ๅๆฅไปปใ", | |
| "labels": ["ไบบๅๆๅงๅ", "ๅ ฌๅธๆ็ป็ปๆบๆๅ็งฐ"], | |
| "language": "zh"}) | |
| texts = {e["text"] for e in resp.json()["entities"]} | |
| assert {"้ฉฌไบ", "ๅผ ๅ", "้ฟ้ๅทดๅทด"} <= texts | |
| def test_chinese_entity_boundary(client): | |
| """BERT NER ๅบ็ฒพ็กฎๆชๆญๅฎไฝ่พน็๏ผไธๅซๅจ่ฏใ""" | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="ๅฐคๆฐ", label="ไบบๅๆๅงๅ", score=0.82, start=0, end=2), | |
| Entity(text="็็ๅค", label="ไบบๅๆๅงๅ", score=0.95, start=8, end=11), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "ๅฐคๆฐๆฅ่ฏท๏ผ็็ๅค็ฌ้๏ผ'ไฝ ๆฅไบใ'", | |
| "labels": ["ไบบๅๆๅงๅ"], "language": "zh"}) | |
| texts = {e["text"] for e in resp.json()["entities"]} | |
| assert "ๅฐคๆฐ" in texts | |
| assert "็็ๅค" in texts | |
| assert "ๅฐคๆฐๆฅ่ฏท" not in texts | |
| assert "็็ๅค็ฌ้" not in texts | |
| def test_chinese_location_product(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="ๆญๅท", label="ๅฐๅๆๅๅธ", score=0.93, start=9, end=11), | |
| Entity(text="ๆทๅฎ", label="ไบงๅๆๅ็ๅ็งฐ", score=0.91, start=14, end=16), | |
| Entity(text="ๅคฉ็ซ", label="ไบงๅๆๅ็ๅ็งฐ", score=0.92, start=17, end=19), | |
| Entity(text="ๆฏไปๅฎ", label="ไบงๅๆๅ็ๅ็งฐ", score=0.90, start=20, end=23), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "้ฟ้ๅทดๅทดๆป้จไฝไบๆญๅท๏ผๆไธๆๆทๅฎใๅคฉ็ซใๆฏไปๅฎใ", | |
| "labels": ["ๅฐๅๆๅๅธ", "ไบงๅๆๅ็ๅ็งฐ"], | |
| "language": "zh"}) | |
| texts = {e["text"] for e in resp.json()["entities"]} | |
| assert {"ๆญๅท", "ๆทๅฎ", "ๅคฉ็ซ", "ๆฏไปๅฎ"} <= texts | |
| def test_chinese_auto_routes_to_zh(client): | |
| """auto ๆฃๆตๅฐไธญๆๅบ่ทฏ็ฑๅฐ ZH ๆจกๅ๏ผlanguage ้ไผ ไธบ 'auto'๏ผๅ ้จๆฃๆตไธบ zh๏ผใ""" | |
| c, mock_ner = client | |
| c.post("/api/v1/extract", | |
| json={"text": "้ฉฌไบๅ็ซไบ้ฟ้ๅทดๅทดใ"}) | |
| # NERService.extract ่ขซ่ฐ็จๆถ language='auto'๏ผ่ทฏ็ฑ้ป่พๅจ ner.py ๅ ้จๅค็ | |
| mock_ner.extract.assert_called_once() | |
| call_kwargs = mock_ner.extract.call_args | |
| assert call_kwargs[1].get("language") == "auto" or call_kwargs[0][3] == "auto" | |
| # โโ Arabic โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_arabic_person_location(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="ู ุญู ุฏ ุจู ุณูู ุงู", label="full name of a person", score=0.82, start=12, end=26), | |
| Entity(text="ุงูู ู ููุฉ ุงูุนุฑุจูุฉ ุงูุณุนูุฏูุฉ", label="geographical location", score=0.85, start=44, end=68), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "ุฃุนูู ุงูุฑุฆูุณ ู ุญู ุฏ ุจู ุณูู ุงู ู ุดุฑูุน ูููู ูู ุงูู ู ููุฉ ุงูุนุฑุจูุฉ ุงูุณุนูุฏูุฉ.", | |
| "labels": ["full name of a person", "geographical location"], | |
| "language": "ar"}) | |
| texts = {e["text"] for e in resp.json()["entities"]} | |
| assert "ู ุญู ุฏ ุจู ุณูู ุงู" in texts | |
| assert "ุงูู ู ููุฉ ุงูุนุฑุจูุฉ ุงูุณุนูุฏูุฉ" in texts | |
| # โโ Mixed Chinese-English โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_mixed_entities_both_scripts(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="ๅผ ไผ", label="person", score=0.95, start=0, end=2), | |
| Entity(text="Google", label="organization", score=0.97, start=9, end=15), | |
| Entity(text="ๅไบฌ", label="location", score=0.93, start=25, end=27), | |
| Entity(text="Android", label="product", score=0.91, start=33, end=40), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "ๅผ ไผๅ ฅ่ไบ Google๏ผ้ฉปๆๅจๅไบฌ๏ผ่ด่ดฃ Android ็ ๅใ", | |
| "labels": ["full name of a person", "ไบบๅๆๅงๅ", | |
| "company or organization name", "ๅ ฌๅธๆ็ป็ปๆบๆๅ็งฐ", | |
| "geographical location", "ๅฐๅๆๅๅธ", | |
| "product or technology name"], | |
| "language": "mixed"}) | |
| texts = {e["text"] for e in resp.json()["entities"]} | |
| assert {"ๅผ ไผ", "Google", "ๅไบฌ", "Android"} <= texts | |
| def test_mixed_no_cross_language_contamination(client): | |
| c, mock_ner = client | |
| mock_ner.extract.return_value = _ents( | |
| Entity(text="OpenAI", label="organization", score=0.97, start=5, end=11), | |
| Entity(text="็่ณ", label="person", score=0.93, start=15, end=17), | |
| ) | |
| resp = c.post("/api/v1/extract", | |
| json={"text": "ไปๅจ OpenAI ๅทฅไฝ๏ผๅไบ็่ณไนๅจๅไธ้จ้จใ", | |
| "labels": ["person", "organization"]}) | |
| entities = resp.json()["entities"] | |
| assert any(e["text"] == "OpenAI" and e["label"] == "organization" for e in entities) | |
| assert any(e["text"] == "็่ณ" and e["label"] == "person" for e in entities) | |
| # โโ Fallback & merge (NERService unit tests, no HTTP) โโโโโโโโโโโโโโโโโโโโโโโโ | |
| def _build_svc(): | |
| """Construct a bare NERService with mocked backends and locks.""" | |
| import threading | |
| from app.ner import NERService | |
| svc = NERService.__new__(NERService) | |
| svc._en_lock = threading.Lock() | |
| svc._zh_lock = threading.Lock() | |
| svc._en_backend = MagicMock() | |
| svc._zh_backend = MagicMock() | |
| return svc | |
| # โโ Sufficiency heuristic โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_expected_min_short_text(): | |
| from app.ner import NERService | |
| assert NERService._expected_min("้ฉฌไบ", []) == 1 | |
| def test_expected_min_medium_text(): | |
| from app.ner import NERService | |
| text = "x" * 50 | |
| assert NERService._expected_min(text, []) == 2 | |
| def test_expected_min_long_text(): | |
| from app.ner import NERService | |
| text = "x" * 350 | |
| assert NERService._expected_min(text, []) == 4 | |
| def test_expected_min_label_floor_takes_over(): | |
| """9 ไธชๆ ็ญพ โ โ9/3โ=3๏ผ่ถ ่ฟ็ญๆๆฌ็ length_floor=1๏ผๆ็ปๅ 3ใ""" | |
| from app.ner import NERService | |
| short_text = "้ฉฌไบ" | |
| labels = [f"l{i}" for i in range(9)] | |
| assert NERService._expected_min(short_text, labels) == 3 | |
| # โโ ZH branch fallback โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_zh_empty_triggers_fallback_and_adds(): | |
| """ZH ไธปๆจกๅ 0 ไธช โ ่งฆๅๅ ๅบ โ ่ฟๅๅ ๅบ็ปๆใ""" | |
| svc = _build_svc() | |
| svc._zh_backend.predict.return_value = ([], []) | |
| svc._en_backend.predict.return_value = _ents( | |
| Entity(text="้ฉฌไบ", label="person", score=0.75, start=0, end=2) | |
| ) | |
| entities, _ = svc.extract("้ฉฌไบ", [], 0.4, language="zh") | |
| assert any(e.text == "้ฉฌไบ" for e in entities) | |
| svc._zh_backend.predict.assert_called_once() | |
| svc._en_backend.predict.assert_called_once() | |
| def test_zh_sufficient_no_fallback(): | |
| """ZH ไธปๆจกๅๅฎไฝๆฐ โฅ expected_min(=1 ็ญๆๆฌ) โ ไธ่ฐ็จๅ ๅบใ""" | |
| svc = _build_svc() | |
| svc._zh_backend.predict.return_value = _ents( | |
| Entity(text="้ฉฌไบ", label="person", score=0.92, start=0, end=2) | |
| ) | |
| svc.extract("้ฉฌไบ", [], 0.4, language="zh") | |
| svc._en_backend.predict.assert_not_called() | |
| def test_zh_insufficient_triggers_fallback_and_results_added(): | |
| """ | |
| ๅ ณ้ฎๆต่ฏ๏ผZH ่ฟๅ 1 ไธช๏ผไฝๆๆฌ้ฟ โ expected_min=4๏ผไธๅ ๅ โ | |
| ่งฆๅๅ ๅบ๏ผไธป็ปๆ + ๅ ๅบ็ปๆไธๅนถ่ฟๅ๏ผ็ธๅ ๏ผไธๆฟๆข๏ผใ | |
| """ | |
| svc = _build_svc() | |
| long_text = "้ฉฌไบ" + "x" * 350 # length_floor = 4 | |
| svc._zh_backend.predict.return_value = _ents( | |
| Entity(text="้ฉฌไบ", label="ไบบๅๆๅงๅ", score=0.95, start=0, end=2), | |
| ) | |
| svc._en_backend.predict.return_value = _ents( | |
| Entity(text="Tesla", label="organization", score=0.90, start=10, end=15), | |
| Entity(text="2024", label="date", score=0.88, start=20, end=24), | |
| ) | |
| entities, _ = svc.extract(long_text, [], 0.4, language="zh") | |
| texts = {e.text for e in entities} | |
| # ไธปๆจกๅ็"้ฉฌไบ"ๅฟ ้กปไฟ็๏ผๅๆถๅ ๅบ็ Tesla / 2024 ไนๅ ่ฟๆฅ | |
| assert "้ฉฌไบ" in texts | |
| assert "Tesla" in texts | |
| assert "2024" in texts | |
| assert len(entities) == 3 # 1 + 2 = 3๏ผ็กฎๅฎๆฏ็ธๅ | |
| def test_user_min_entities_overrides_heuristic(): | |
| """่ฏทๆฑ้ไผ min_entities=5 ๆถๅบ่ฆ็ๅฏๅๅผ๏ผไธปๆจกๅ 3 ไธชไป่งฆๅๅ ๅบใ""" | |
| svc = _build_svc() | |
| svc._zh_backend.predict.return_value = _ents( | |
| Entity(text="้ฉฌไบ", label="person", score=0.95, start=0, end=2), | |
| Entity(text="ๅผ ๅ", label="person", score=0.93, start=4, end=6), | |
| Entity(text="ๆญๅท", label="location", score=0.91, start=8, end=10), | |
| ) | |
| svc._en_backend.predict.return_value = _ents( | |
| Entity(text="Tesla", label="organization", score=0.85, start=15, end=20), | |
| ) | |
| entities, _ = svc.extract("้ฉฌไบใๅผ ๅใๆญๅท Tesla", [], 0.4, | |
| language="zh", min_entities=5) | |
| # 3 < 5 โ ่งฆๅๅ ๅบ๏ผๆ็ป 3 + 1 = 4 | |
| assert len(entities) == 4 | |
| svc._en_backend.predict.assert_called_once() | |
| def test_user_min_entities_zero_disables_fallback(): | |
| """min_entities=0 ๆถไธปๆจกๅๅณไฝฟ่ฟๅ็ฉบไนไธ่งฆๅๅ ๅบใ""" | |
| svc = _build_svc() | |
| svc._zh_backend.predict.return_value = ([], []) | |
| svc.extract("้ฉฌไบ", [], 0.4, language="zh", min_entities=0) | |
| svc._en_backend.predict.assert_not_called() | |
| # โโ EN branch fallback (symmetric) โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_en_insufficient_triggers_zh_fallback_and_adds(): | |
| """EN ไธปๆจกๅไธๅ ๅ โ ่ฐ ZH ๅ ๅบ โ ็ปๆ็ธๅ ใ""" | |
| svc = _build_svc() | |
| long_text = "Tesla" + "x" * 350 | |
| svc._en_backend.predict.return_value = _ents( | |
| Entity(text="Tesla", label="organization", score=0.95, start=0, end=5), | |
| ) | |
| svc._zh_backend.predict.return_value = _ents( | |
| Entity(text="้ฉฌไบ", label="ไบบๅๆๅงๅ", score=0.91, start=10, end=12), | |
| ) | |
| entities, _ = svc.extract(long_text, [], 0.4, language="en") | |
| texts = {e.text for e in entities} | |
| assert "Tesla" in texts and "้ฉฌไบ" in texts | |
| assert len(entities) == 2 | |
| # โโ Mixed: always merge both โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ | |
| def test_mixed_always_runs_both_models(): | |
| """mixed ่ฏญ่จๆ ่งๅ ๅๆง๏ผๆฐธ่ฟ่ทไธคไธชๆจกๅๅนถๅๅนถใ""" | |
| svc = _build_svc() | |
| svc._en_backend.predict.return_value = _ents( | |
| Entity(text="Google", label="organization", score=0.95, start=5, end=11) | |
| ) | |
| svc._zh_backend.predict.return_value = _ents( | |
| Entity(text="ๅผ ไผ", label="ไบบๅๆๅงๅ", score=0.91, start=0, end=2) | |
| ) | |
| entities, _ = svc.extract("ๅผ ไผๅ ๅ ฅ Googleใ", [], 0.4, language="mixed") | |
| texts = {e.text for e in entities} | |
| assert {"Google", "ๅผ ไผ"} <= texts | |
| svc._en_backend.predict.assert_called_once() | |
| svc._zh_backend.predict.assert_called_once() | |
| def test_merge_deduplicates_overlapping_spans(): | |
| """ไธคไธชๆจกๅๅฏนๅไธ span ้ฝๅฝไธญ โ ไฟ็ๅพๅๆ้ซ็้ฃๆกใ""" | |
| svc = _build_svc() | |
| svc._en_backend.predict.return_value = ( | |
| [Entity(text="ๅผ ไผ", label="person", score=0.70, start=0, end=2)], | |
| ["person"], | |
| ) | |
| svc._zh_backend.predict.return_value = ( | |
| [Entity(text="ๅผ ไผ", label="ไบบๅๆๅงๅ", score=0.92, start=0, end=2)], | |
| ["ไบบๅๆๅงๅ"], | |
| ) | |
| entities, _ = svc.extract("ๅผ ไผ", [], 0.4, language="mixed") | |
| matches = [e for e in entities if e.text == "ๅผ ไผ"] | |
| assert len(matches) == 1 | |
| assert matches[0].score == 0.92 # ้ซๅ่ๅบ | |