-
Notifications
You must be signed in to change notification settings - Fork 1
/
heideltime_wrapper.py
49 lines (36 loc) · 1.82 KB
/
heideltime_wrapper.py
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
"""
Wrapper for Heideltime.
"""
from typing import Tuple, Dict, Optional
from functools import lru_cache
from python_heideltime import Heideltime
heideltime_cache = {"heideltime": Heideltime()}
LANGUAGE_MAPPING = {"en": "ENGLISH", "de": "GERMAN", "es": "SPANISH", "it": "ITALIAN", "fr": "FRENCH",
"nl": "DUTCH", "et": "ESTONIAN", "pt": "PORTUGUESE"}
HEIDELTIME_LANGUAGES = set(LANGUAGE_MAPPING.keys())
@lru_cache(maxsize=24)
def heideltime_prediction(texts: Tuple[str], language: str, heideltime_mode: str, date: Optional[str]) \
-> Tuple[Dict, int]:
# Special language considerations for English colloquial and scientific texts
if heideltime_mode == "colloquial" and language == "en":
heideltime_language_tag = "ENGLISHCOLL"
elif heideltime_mode == "scientific" and language == "en":
heideltime_language_tag = "ENGLISHSCI"
else:
heideltime_language_tag = determine_language(language)
# Set document parameters
heideltime_cache["heideltime"].set_document_type(heideltime_mode.upper())
heideltime_cache["heideltime"].set_language(heideltime_language_tag)
if date:
heideltime_cache["heideltime"].set_document_time(date)
# Join texts back together before processing, as Heideltime can use this to cross-reference date anchors.
# Also, spinning up Heideltime takes some overhead time, so reducing the number of calls is advantageous, too.
doc = heideltime_cache["heideltime"].parse('\n'.join(texts))
doc = strip_xml_wrapper(doc)
return {"tagged_text": doc}, 200
def determine_language(lang: str) -> str:
return LANGUAGE_MAPPING[lang]
def strip_xml_wrapper(text: str) -> str:
text = text.replace('<?xml version="1.0"?>\n<!DOCTYPE TimeML SYSTEM "TimeML.dtd">\n<TimeML>\n', '')
text = text.replace('\n</TimeML>\n\n', '')
return text