Spaces:
Build error
Build error
| #modules/extractor_module.py | |
| from typing import Dict, List | |
| from core.base_module import AIModule | |
| import spacy | |
| class ExtractorModule(AIModule): | |
| def __init__(self): | |
| self.nlp = spacy.load("en_core_web_sm") | |
| async def process(self, input_data: Dict) -> Dict: | |
| text = input_data.get("text", "") | |
| doc = self.nlp(text) | |
| entities = [ | |
| {"text": ent.text, "label": ent.label_} | |
| for ent in doc.ents | |
| ] | |
| return { | |
| "entities": entities, | |
| "num_entities": len(entities) | |
| } | |
| async def get_status(self) -> Dict: | |
| return {"status": "operational", "model": "en_core_web_sm"} | |
| def capabilities(self) -> List[str]: | |
| return ["named-entity-recognition", "information-extraction"] | |