import re from os import path from collections import Counter from loguru import logger # from langdetect import detect import spacy import en_core_web_sm import zh_core_web_sm from magic_pdf.libs.language import detect_lang class NLPModels: """ How to upload local models to s3: - config aws cli: doc\SETUP-CLI.md doc\setup_cli.sh app\config\__init__.py - $ cd {local_dir_storing_models} - $ ls models en_core_web_sm-3.7.1/ zh_core_web_sm-3.7.0/ - $ aws s3 sync models/ s3://llm-infra/models --profile=p_project_norm - $ aws s3 --profile=p_project_norm ls s3://llm-infra/models/ PRE en_core_web_sm-3.7.1/ PRE zh_core_web_sm-3.7.0/ """ def __init__(self): # if OS is windows, set "TMP_DIR" to "D:/tmp" home_dir = path.expanduser("~") self.default_local_path = path.join(home_dir, ".nlp_models") self.default_shared_path = "/share/pdf_processor/nlp_models" self.default_hdfs_path = "hdfs://pdf_processor/nlp_models" self.default_s3_path = "s3://llm-infra/models" self.nlp_models = self.nlp_models = { "en_core_web_sm": { "type": "spacy", "version": "3.7.1", }, "en_core_web_md": { "type": "spacy", "version": "3.7.1", }, "en_core_web_lg": { "type": "spacy", "version": "3.7.1", }, "zh_core_web_sm": { "type": "spacy", "version": "3.7.0", }, "zh_core_web_md": { "type": "spacy", "version": "3.7.0", }, "zh_core_web_lg": { "type": "spacy", "version": "3.7.0", }, } self.en_core_web_sm_model = en_core_web_sm.load() self.zh_core_web_sm_model = zh_core_web_sm.load() def load_model(self, model_name, model_type, model_version): if ( model_name in self.nlp_models and self.nlp_models[model_name]["type"] == model_type and self.nlp_models[model_name]["version"] == model_version ): return spacy.load(model_name) if spacy.util.is_package(model_name) else None else: logger.error(f"Unsupported model name or version: {model_name} {model_version}") return None def detect_language(self, text, use_langdetect=False): if len(text) == 0: return None if use_langdetect: # print("use_langdetect") # print(detect_lang(text)) # return detect_lang(text) if detect_lang(text) == "zh": return "zh" else: return "en" if not use_langdetect: en_count = len(re.findall(r"[a-zA-Z]", text)) cn_count = len(re.findall(r"[\u4e00-\u9fff]", text)) if en_count > cn_count: return "en" if cn_count > en_count: return "zh" def detect_entity_catgr_using_nlp(self, text, threshold=0.5): """ Detect entity categories using NLP models and return the most frequent entity types. Parameters ---------- text : str Text to be processed. Returns ------- str The most frequent entity type. """ lang = self.detect_language(text, use_langdetect=True) if lang == "en": nlp_model = self.en_core_web_sm_model elif lang == "zh": nlp_model = self.zh_core_web_sm_model else: # logger.error(f"Unsupported language: {lang}") return {} # Splitting text into smaller parts text_parts = re.split(r"[,;，；、\s & |]+", text) text_parts = [part for part in text_parts if not re.match(r"[\d\W]+", part)] # Remove non-words text_combined = " ".join(text_parts) try: doc = nlp_model(text_combined) entity_counts = Counter([ent.label_ for ent in doc.ents]) word_counts_in_entities = Counter() for ent in doc.ents: word_counts_in_entities[ent.label_] += len(ent.text.split()) total_words_in_entities = sum(word_counts_in_entities.values()) total_words = len([token for token in doc if not token.is_punct]) if total_words_in_entities == 0 or total_words == 0: return None entity_percentage = total_words_in_entities / total_words if entity_percentage < 0.5: return None most_common_entity, word_count = word_counts_in_entities.most_common(1)[0] entity_percentage = word_count / total_words_in_entities if entity_percentage >= threshold: return most_common_entity else: return None except Exception as e: logger.error(f"Error in entity detection: {e}") return None def __main__(): nlpModel = NLPModels() test_strings = [ "张三", "张三, 李四，王五; 赵六", "John Doe", "Jane Smith", "Lee, John", "John Doe, Jane Smith; Alice Johnson，Bob Lee", "孙七, Michael Jordan；赵八", "David Smith Michael O'Connor; Kevin ßáçøñ", "李雷·韩梅梅, 张三·李四", "Charles Robert Darwin, Isaac Newton", "莱昂纳多·迪卡普里奥, 杰克·吉伦哈尔", "John Doe, Jane Smith; Alice Johnson", "张三, 李四，王五; 赵六", "Lei Wang, Jia Li, and Xiaojun Chen, LINKE YANG OU, and YUAN ZHANG", "Rachel Mills & William Barry & Susanne B. Haga", "Claire Chabut* and Jean-François Bussières", "1 Department of Chemistry, Northeastern University, Shenyang 110004, China 2 State Key Laboratory of Polymer Physics and Chemistry, Changchun Institute of Applied Chemistry, Chinese Academy of Sciences, Changchun 130022, China", "Changchun", "china", "Rongjun Song, 1,2 Baoyan Zhang, 1 Baotong Huang, 2 Tao Tang 2", "Synergistic Effect of Supported Nickel Catalyst with Intumescent Flame-Retardants on Flame Retardancy and Thermal Stability of Polypropylene", "Synergistic Effect of Supported Nickel Catalyst with", "Intumescent Flame-Retardants on Flame Retardancy", "and Thermal Stability of Polypropylene", ] for test in test_strings: print() print(f"Original String: {test}") result = nlpModel.detect_entity_catgr_using_nlp(test) print(f"Detected entities: {result}") if __name__ == "__main__": __main__()