BytecodeApps/docverse-api

1import spacy2from typing import Dict, List, Any3 4# NLP Pipeline using SciSpacy (Requires: en_core_sci_sm or similar model)5# pip install scispacy6# pip install https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.5.1/en_core_sci_sm-0.5.1.tar.gz7 8class NLPParser:9    def __init__(self, model_name="en_core_sci_sm"):10        try:11            self.nlp = spacy.load(model_name)12        except OSError:13            # Fallback for demonstration if scispacy is not installed locally14            print(f"Warning: {model_name} not found. Attempting to load 'en_core_web_sm'.")15            try:16                self.nlp = spacy.load("en_core_web_sm")17            except OSError:18                print("Warning: No Spacy models found. Using mock parser for demonstration.")19                self.nlp = None20 21    def extract_entities(self, text: str) -> Dict[str, List[Any]]:22        """23        Extracts entities such as medications, conditions, and lab values from raw text.24        """25        if not self.nlp:26            # Return mock data if model failed to load (for local Dev/Testing without ML dependencies)27            return self._mock_extraction(text)28        29        doc = self.nlp(text)30        31        entities = {32            "disease_entity": [],33            "chemical_entity": [], # Medications34            "other_entities": []35        }36 37        # SciSpacy uses specific labels like 'ENTITY', 'DISEASE', 'CHEMICAL' depending on the model38        for ent in doc.ents:39            # Pseudo-logic: SciSpacy base models usually tag generic 'ENTITY'40            # Specialized models might tag 'CHEMICAL', 'DISEASE'41            label = ent.label_42            if label in ['CHEMICAL', 'DRUG']:43                entities["chemical_entity"].append(ent.text)44            elif label in ['DISEASE', 'CONDITION']:45                entities["disease_entity"].append(ent.text)46            else:47                entities["other_entities"].append({"text": ent.text, "label": label})48        49        # Deduplicate50        entities["chemical_entity"] = list(set(entities["chemical_entity"]))51        entities["disease_entity"] = list(set(entities["disease_entity"]))52        53        return entities54 55    def _mock_extraction(self, text: str):56        """ Fallback logic if spacy isn't installed during build phase """57        results = {58            "chemical_entity": [],59            "disease_entity": [],60            "extracted_vitals": []61        }62        text_lower = text.lower()63        if "lisinopril" in text_lower: results["chemical_entity"].append("Lisinopril")64        if "sildenafil" in text_lower: results["chemical_entity"].append("Sildenafil")65        if "hemoglobin" in text_lower: results["extracted_vitals"].append({"test": "Hemoglobin", "value": 6.5, "unit": "g/dL"})66        if "potassium" in text_lower: results["extracted_vitals"].append({"test": "Potassium", "value": 6.2, "unit": "mEq/L"})67        if "hypertension" in text_lower: results["disease_entity"].append("Hypertension")68        return results69 70# -- Usage Example --71# parser = NLPParser()72# entities = parser.extract_entities("Patient was prescribed Lisinopril for Hypertension. Hemoglobin is 6.5 g/dL.")73