ZeroShot Multitask - CDA/XML PHI Deidentification

from johnsnowlabs import nlp, medical
document_assembler = nlp.DocumentAssembler()\
    .setInputCol("text")\
    .setOutputCol("document")

zero_shot = medical.PretrainedZeroShotMultiTask.pretrained("zeroshot_multitask_deid_cda", "en", "clinical/models")\
    .setInputCols(["document"])\
    .setOutputCol("extractions")\
    .setEntityThreshold(0.4)\
    .setEntities([
        "PATIENT_NAME::str::All mentions of person names, including patients, providers, family members, or other individuals named in the record. Examples: 'Alice Jones', 'Dr. Smith', 'John Doe', etc.",
        "DATE_or_Time::str::Any mention of dates, including birth dates, admission dates, and other chronological references. Examples: '05/01/1970', 'May 1, 1970', '2020-05-01', '5/1/70', 'May 2020', '2020', etc.",
        "LOCATION::str::Any mention of street addresses, cities, states, postal codes, or other location identifiers. Examples: '123 Main St', 'Apt 4B', 'New York, NY 10001', etc.",
        "TELECOM::str::Any mention of telephone numbers in any format, including extensions and international codes. Examples: '+1(555)-777-1234', '(555)-723-1544'",
    ])

pipeline = nlp.Pipeline(
    stages = [
        document_assembler,
        zero_shot
])

text = f"""
<recordTarget xmlns="urn:hl7-org:v3" xmlns:sdtc="urn:hl7-org:sdtc">
    <patientRole>
      <id extension="T-10118" root="2.16.840.1.113883.4.1" />
      <addr use="HP">
        <streetAddressLine>1357 Amber Dr</streetAddressLine>
        <city>Beaverton</city>
        <state>OR</state>
        <postalCode>97006</postalCode>
        <country>US</country>
      </addr>
      <telecom use="MC" value="tel:+1(555)-777-1234" />
      <telecom use="HP" value="tel:+1(555)-723-1544" />
      <telecom value="360mu.alice.newman@gmail.com" />
      <patient>
        <name use="L">
          <given>Alice</given>
          <given>Jones</given>
          <family>Newman</family>
        </name>
        <name>
          <given qualifier="BR">Alicia</given>
          <family>Newman</family>
        </name>
        <administrativeGenderCode code="F" codeSystem="2.16.840.1.113883.5.1" codeSystemName="AdministrativeGender" displayName="Female" />
        <birthTime value="19700501" />
        
    
<entry xmlns="urn:hl7-org:v3" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" typeCode="DRIV">
            <act classCode="ACT" moodCode="EVN">
              <templateId extension="2015-08-01" root="2.16.840.1.113883.10.20.22.4.30" />
              <templateId root="2.16.840.1.113883.10.20.22.4.30" />
              <id root="b03805bd-2eb6-4ab8-a9ff-473c6653971a" />
              <code code="CONC" codeSystem="2.16.840.1.113883.5.6" />
              <statusCode code="active" />
              <effectiveTime>
                <low value="19800510" />
              </effectiveTime>
              <entryRelationship typeCode="SUBJ">
                <observation classCode="OBS" moodCode="EVN">
                  <templateId extension="2014-06-09" root="2.16.840.1.113883.10.20.22.4.7" />
                  <templateId root="2.16.840.1.113883.10.20.22.4.7" />
                  <id root="4adc1020-7b14-11db-9fe1-0800200c9a68" />
                  <code code="ASSERTION" codeSystem="2.16.840.1.113883.5.4" />
                  <statusCode code="completed" />
                  <effectiveTime>
                    <low value="19800510" />
                  </effectiveTime>
                  <value code="419511003" codeSystem="2.16.840.1.113883.6.96" codeSystemName="SNOMED CT" displayName="Propensity to adverse reaction to drug" xsi:type="CD" />
                  <participant typeCode="CSM">
                    <participantRole classCode="MANU">
                      <playingEntity classCode="MMAT">
                        <code code="733" codeSystem="2.16.840.1.113883.6.88" codeSystemName="RxNorm" displayName="Ampicillin">
                          <originalText>
                            <reference value="#product2" />
                          </originalText>
                        </code>
                      </playingEntity>
                    </participantRole>
                  </participant>
                  <entryRelationship inversionInd="true" typeCode="MFST">
                    <observation classCode="OBS" moodCode="EVN">
                      <templateId extension="2014-06-09" root="2.16.840.1.113883.10.20.22.4.9" />
                      <templateId root="2.16.840.1.113883.10.20.22.4.9" />
                      <id root="4adc1020-7b14-11db-9fe1-0800200c9a69" />
                      <code code="ASSERTION" codeSystem="2.16.840.1.113883.5.4" />
                      <text>
                        <reference value="#reaction2" />
                      </text>
                      <statusCode code="completed" />
                      <effectiveTime>
                        <low nullFlavor="NI" />
                        <high nullFlavor="NI" />
                      </effectiveTime>
                      <value code="247472004" codeSystem="2.16.840.1.113883.6.96" displayName="Hives" xsi:type="CD" />
                      <entryRelationship inversionInd="true" typeCode="SUBJ">
                        <observation classCode="OBS" moodCode="EVN">
"""

data = spark.createDataFrame([[text]]).toDF("text")

results = pipeline.fit(data).transform(data)

results.select("extractions").show(truncate=False)