<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Read the parquet file into a pandas dataframe. Using fastparquet here because pyarrow couldn't read in a file of this size for some reason
Step1:  Get the list of ids from the processed xml files so we can select a subset of the mimic notes
Step2:  Select the subset of notes that we have xml output from ctakes for. Reset the index and drop some unnecessary columns
Step3:  process the xml files and store in parquet locally
TODO
Step4:  Creating templates
The plan
Step5:  Pull in the dataframes for elements we need for processing
Step6:  Prep sentences DF
Add raw text from notes to sentences
Step7:  Add position of sentence in document to sentences df
Step8:  remove sentences without entities
Step9:  Prep UMLS DF
Remove umls concepts which don't have a preferred text field
Step10:  Pref Mentions DF
Transform mentions
Drop some unused fields
Only keep the first umls code from ontology array ( no longer doing this as it limits the cui codes we can choose
    from in the umls concepts table)
Sort by begin and end offsets.  Remove mentions that end on the same offset.  Only want to keep the full span and not split entities up.  This should give us better semantic meaning
Add raw text to mentions
Add in umls concept information (CUI) to mentions
    a. There are many possible cuis for a the text span of an entity.  Here, we're going to use 
    the edit distance from the original span and the umls preferred text.  For now, just choose the first 
    umls concept with the best score (lowest)
Step11:  Add original text to mentions
Step12:  Add sentence position to mentions
Step13:  Prep Predicates DF
Transform predicates
Simple transformation. Just modify the frameset string to remove everything after the '.'
Step14:  Remove predicates not in sentences with mentions
Step15:  Add original text to predicates
Step16:  Linking CUI codes to entities (mentions)
Assign cui codes to mentions (entities)
cTAKES over-generates cui and tui codes for text spans in a clinical note.  There can be multiple coding schemes that have a code for a term and a cui could apply to the original text span specifically or be a generalization or abstraction over the meaning of the span.  For generating text we want the cui that most closely matches the original text span.  Future work could look at these generalizations to get a better sense of semantic meaning.  However, this will require a deep understanding of the UMLS ontology an how to work with it to extract this kind of information.
For each mention
Step17:  Set the template tokens we're going to use
For mentions this is either
Step18:  Append the two template tokens dataframes
Step19:  Get the semantic templates
Group the rows of the above template tokens dataframe by sentence id and join them together into a single string.  Must sort by begin offset.
Step20:  Gather corpus statistics
Average sentences per doc
Step21:  Average sentences w/ entities per doc
Step22:  Count of unique cuis (When removing overlapping text spans)
Step23:  Average # of cuis per doc
Step24:  Average # of cuis per sentence
Step25:  Average # of words per doc (excluding newline tokens and symbols)
Step26:  Average # of words per sentence
Step27:  Get frequency of mentions
Step28:  Frequency of mentions by sentence position
Step29:  Frequency of CUIs
Step30:  Frequency with preferred text
Step31:  Frequency of CUIs by sentence position
Step32:  Number of unique templates
Step33:  Frequency of templates (identified by sentence number)
Step34:  Frequency of templates by sentence position
Step35:  Write dataframes to parquet
We want to write these to a parquet file so that they can be used by a separate notebook to do clustering and note generation. This is just prep-work for those processes.
<ASSISTANT_TASK:>
Python Code:
notes_file = 'synthnotes/data/note-events.parquet'
pq_root_path = 'synthnotes/data/xml_extracted'
pf = ParquetFile(notes_file)
df = pf.to_pandas()
Explanation: Read the parquet file into a pandas dataframe. Using fastparquet here because pyarrow couldn't read in a file of this size for some reason
End of explanation
xml_dir = 'synthnotes/data/xml_files'
xml_files = os.listdir(xml_dir)
ids = [int(f.split('.txt.xmi')[0]) for f in xml_files]
Explanation: Get the list of ids from the processed xml files so we can select a subset of the mimic notes
End of explanation
notes = df[df.ROW_ID.isin(ids)]
notes = notes.reset_index(drop=True)
notes = notes.drop(['CHARTDATE','CHARTTIME','STORETIME','CGID','ISERROR'],axis=1)
def get_notes_sample(df, n=100, category='Nursing'):
    notes = df[notes_df['CATEGORY'] == 'Nursing']
    notes = notes[notes['ISERROR'].isnull()]
    notes = notes[notes['DESCRIPTION'] == 'Generic Note']
    notes = notes.sample(n=n)
    notes = notes.reset_index(drop=True)
    return notes
    
Explanation: Select the subset of notes that we have xml output from ctakes for. Reset the index and drop some unnecessary columns
End of explanation
# parser = CtakesXmlParser()
# schemas = list()
# for file in xml_files:
#     xml_out = parser.parse(f'{xml_dir}/{file}')
#     for k, v in xml_out.items():
#         feature_df = pd.DataFrame(list(v))  
#         if feature_df.shape[0] > 0:
#             table = pa.Table.from_pandas(feature_df)
#             pq.write_to_dataset(table, f'{pq_root_path}/{k}')
#         else:
#             print(f"{k} was empty for {file}")
Explanation: process the xml files and store in parquet locally
TODO: switch this to use columnar format: need to change how we extract different types of elements
End of explanation
def get_df_from_pq(root, name):
    return pq.read_table(f'{root}/{name}').to_pandas()
def transform_preds(df):
    df['frameset'] = df['frameset'].apply(lambda x: x.split('.')[0])
    return df
def transform_mentions(mentions):
    # Don't want this to fail if these have already been removed
    try:
        mentions = mentions.drop(
            ['conditional', 'history_of', 'generic', 'polarity', 'discovery_technique', 'subject'],
            axis=1)
    except:
        pass
    
    sorted_df = mentions.groupby(['sent_id', 'begin']) \
                        .apply(lambda x: x.sort_values(['begin', 'end']))
    
    # Drop the mentions that are parts of a larger span.  Only keep the containing span that holds multiple
    # mentions
    deduped = sorted_df.drop_duplicates(['sent_id', 'begin'], keep='last')
    deduped = deduped.drop_duplicates(['sent_id', 'end'], keep='first')
    return deduped.reset_index(drop=True)
def set_template_token(df, column):
    df['template_token'] = df[column]
    return df
def get_template_tokens(row):
    return pd.Series({
        'doc_id': row['doc_id'],
        'sent_id': row['sent_id'],
        'token': row['template_token'],
        'begin': row['begin'],
        'end': row['end']
        })    
# def merge_mentions_umls(mentions, umls):
#     umls['umls_xmi_id'] = umls['xmi_id']
#     mentions = mentions.merge(umls[['umls_xmi_id', 'cui']], on='umls_xmi_id')
#     return mentions
# def umls_dedup(umls):
#     return umls.drop_duplicates(subset=['cui'])
# def set_umls_join_key(umls):
#     umls['umls_xmi_id'] = umls['xmi_id']
#     return umls
def set_sentence_pos(df):
    df = df.groupby(["doc_id"]).apply(lambda x: x.sort_values(["begin"])).reset_index(drop=True)
    df['sentence_number'] = df.groupby("doc_id").cumcount()
    return df
def get_root_verb(row):
    pass
def extract_sent(row):
    begin = row['begin']
    end = row['end']
    row['TEXT'] = row['TEXT'][begin:end]
    return row
def write_notes(row):
    fn = f'raw_notes/{row["ROW_ID"]}'
    with open(fn, 'w') as f:
        f.write(row['TEXT'])
        
def get_text_from_sentence(row, notes):
    doc = notes[notes['ROW_ID'] == row['doc_id']]
    b = row['begin']
    e = row['end']
    return doc['TEXT'].iloc[0][b:e]        
def edit_dist(row, term2):
    term1 = row.loc['preferred_text']
    return lev_norm(term1, term2)
    
def get_cui( mention, umls_df):
        ont_arr = list(map(int, mention['ontology_arr'].split())) or None
        ment_text = mention['text']
        
        concepts = umls_df[umls_df['xmi_id'].isin(ont_arr)].loc[:, ['cui', 'preferred_text', 'xmi_id']]
        concepts['dist'] = concepts.apply(edit_dist, args=(ment_text,), axis=1)
        sorted_df = concepts.sort_values(by='dist', ascending=True).reset_index(drop=True)
        cui = sorted_df['cui'].iloc[0]
        xmi_id = sorted_df['xmi_id'].iloc[0]
        pref_text = sorted_df['preferred_text'].iloc[0]
        return cui, xmi_id, pref_text
Explanation: Creating templates
The plan:<br>
For each sentences in all documents:
    1. Get the predicates for the sentence
    2. Get the entities for the sentence
    3. For each entity:
         -  append the cui code from umls concept to the end
    4. Combine predicates and entities and sort based on their begin position
    5. Save to a dataframe
Some helper functions:
End of explanation
preds = get_df_from_pq(pq_root_path, 'predicates')
mentions = get_df_from_pq(pq_root_path, 'mentions')
umls = get_df_from_pq(pq_root_path, 'umls_concepts')
sents = get_df_from_pq(pq_root_path, 'sentences')
tokens = get_df_from_pq(pq_root_path, 'tokens')
sents = sents.rename({'id': 'sent_id'}, axis=1)
sents.head()
Explanation: Pull in the dataframes for elements we need for processing
End of explanation
sents = sents.rename({'id': 'sent_id'}, axis=1)
sents = sents.merge(notes[['ROW_ID', 'TEXT']],
            left_on='doc_id', right_on='ROW_ID').drop('ROW_ID', axis=1)
sents = sents.apply(extract_sent, axis=1)
sents = sents.rename({'TEXT': 'text'}, axis=1)
Explanation: Prep sentences DF
Add raw text from notes to sentences
End of explanation
sents = set_sentence_pos(sents)
Explanation: Add position of sentence in document to sentences df
End of explanation
sents_with_mentions = sents[
    sents['sent_id'].isin(
        mentions.drop_duplicates(subset='sent_id')['sent_id']
    )
]
Explanation: remove sentences without entities
End of explanation
umls = umls[~umls['preferred_text'].isna()]
Explanation: Prep UMLS DF
Remove umls concepts which don't have a preferred text field
End of explanation
mentions = get_df_from_pq(pq_root_path, 'mentions')
mentions = transform_mentions(mentions)
mentions.head()
Explanation: Pref Mentions DF
Transform mentions
Drop some unused fields
Only keep the first umls code from ontology array ( no longer doing this as it limits the cui codes we can choose
    from in the umls concepts table)
Sort by begin and end offsets.  Remove mentions that end on the same offset.  Only want to keep the full span and not split entities up.  This should give us better semantic meaning
Add raw text to mentions
Add in umls concept information (CUI) to mentions
    a. There are many possible cuis for a the text span of an entity.  Here, we're going to use 
    the edit distance from the original span and the umls preferred text.  For now, just choose the first 
    umls concept with the best score (lowest)
End of explanation
mentions['text'] = mentions.apply(get_text_from_sentence, args=(notes,), axis=1)
mentions.head()
Explanation: Add original text to mentions
End of explanation
mentions = mentions.merge(sents_with_mentions[['sent_id', 'sentence_number']],
              on='sent_id')
mentions.head()
Explanation: Add sentence position to mentions
End of explanation
preds = transform_preds(preds)
Explanation: Prep Predicates DF
Transform predicates
Simple transformation. Just modify the frameset string to remove everything after the '.'
End of explanation
print(preds.shape)
preds = preds[
    preds['sent_id'].isin( sents_with_mentions['sent_id'] )
]
print(preds.shape)
Explanation: Remove predicates not in sentences with mentions
End of explanation
preds['text'] = preds.apply(get_text_from_sentence, args=(notes,), axis=1)
Explanation: Add original text to predicates
End of explanation
mentions[['cui', 'umls_xmi_id', 'preferred_text']] = mentions. \
                                                     apply(get_cui, args=(umls,), axis=1, result_type='expand')
mentions.head()
Explanation: Linking CUI codes to entities (mentions)
Assign cui codes to mentions (entities)
cTAKES over-generates cui and tui codes for text spans in a clinical note.  There can be multiple coding schemes that have a code for a term and a cui could apply to the original text span specifically or be a generalization or abstraction over the meaning of the span.  For generating text we want the cui that most closely matches the original text span.  Future work could look at these generalizations to get a better sense of semantic meaning.  However, this will require a deep understanding of the UMLS ontology an how to work with it to extract this kind of information.
For each mention:
1. Collect all the umls concept rows (based on xmi_id) that are in the mention's ontology array
2. Compute edit distance between the above umls rows' preferred text column and the mention's original text
3. Sort edit distances in ascending order
4. Choose the first umls concept row (a lower edit distance means the two texts are more similar)
End of explanation
mentions['template_token'] = mentions['mention_type']
preds['template_token'] = preds['frameset']
preds_toks = preds.apply(get_template_tokens, axis=1)
mentions_toks = mentions.apply(get_template_tokens, axis=1)
mentions_toks.groupby(['sent_id', 'end']).head()
preds_toks.groupby(['sent_id', 'end']).head()
Explanation: Set the template tokens we're going to use
For mentions this is either: the type of mention, the CUI code, or the two concatenated together
For predicates it is the frameset trimmed of everything after the '.'
End of explanation
template_tokens = preds_toks.append(mentions_toks)
temp_tokens = template_tokens.groupby(['sent_id']).apply(lambda x: x.sort_values(['begin']))
temp_tokens.head()
Explanation: Append the two template tokens dataframes
End of explanation
sem_templates = template_tokens.sort_values('begin').groupby('sent_id')['token'].apply(' '.join)
sem_templates.head()
temp_tokens.token.unique().shape
sem_df = pd.DataFrame(sem_templates)  # What is this?
sem_df.head()
sem_df.reset_index(level=0, inplace=True)
sem_df = sem_df.rename(columns={'token': 'sem_template'})
sem_df = sem_df.merge(sents[['sent_id', 'sentence_number', 'doc_id', 'begin', 'end']],
                      left_on='sent_id', right_on='sent_id' )#.drop('id', axis=1)
sem_df.head()
Explanation: Get the semantic templates
Group the rows of the above template tokens dataframe by sentence id and join them together into a single string.  Must sort by begin offset.
End of explanation
avg_sents_per_doc = sents.groupby('doc_id').size().mean()
print(avg_sents_per_doc)
Explanation: Gather corpus statistics
Average sentences per doc
End of explanation
avg_sents_with_ents_per_doc = sents_with_mentions.groupby('doc_id').size().mean()
print(avg_sents_with_ents_per_doc)
Explanation: Average sentences w/ entities per doc
End of explanation
print(mentions['cui'].nunique())
Explanation: Count of unique cuis (When removing overlapping text spans)
End of explanation
mentions.groupby('doc_id').size().mean()
Explanation: Average # of cuis per doc
End of explanation
mentions.groupby('sent_id').size().mean()
Explanation: Average # of cuis per sentence
End of explanation
tokens = tokens[(~tokens['sent_id'].isnull()) & (tokens['token_type'] != 'NewlineToken')]
wc_by_doc = tokens.groupby('doc_id').count()['id'].reset_index(name='count')
wc_by_doc['count'].mean()
Explanation: Average # of words per doc (excluding newline tokens and symbols)
End of explanation
wc_by_sentence = tokens.groupby('sent_id')['id'].count().reset_index(name='count')
wc_by_sentence['count'].mean()
Explanation: Average # of words per sentence
End of explanation
mention_counts = mentions.groupby('mention_type').size().reset_index(name='count')
mention_counts
mention_counts['frequency'] = mention_counts['count'] / mention_counts['count'].sum()
mention_counts
Explanation: Get frequency of mentions
End of explanation
mentions_by_pos = pd.crosstab(
                        mentions['mention_type'],
                        mentions['sentence_number']).apply(lambda x: x / x.sum(), axis=0)
mentions_by_pos
Explanation: Frequency of mentions by sentence position
End of explanation
cui_counts = mentions.groupby('cui').size().reset_index(name='count')
cui_counts = cui_counts.sort_values('count', ascending=False).reset_index(drop=True)
cui_counts.head(10)
cui_counts['frequency'] = cui_counts['count'] / cui_counts['count'].sum()
cui_counts.head(10)
Explanation: Frequency of CUIs
End of explanation
cui_counts_with_text = cui_counts.merge(mentions[['cui', 'preferred_text']], on='cui') \
                        .drop_duplicates('cui') \
                        .reset_index(drop=True)
cui_counts_with_text.head(10)
Explanation: Frequency with preferred text
End of explanation
cui_by_pos = pd.crosstab(mentions['cui'], mentions['sentence_number']).apply(lambda x: x / x.sum(), axis=0)
cui_by_pos.head()
cui_by_pos.loc[:, 0].sort_values(ascending=False)[:10]
Explanation: Frequency of CUIs by sentence position
End of explanation
sem_df.head()
sem_df['sem_template'].nunique()
Explanation: Number of unique templates
End of explanation
count_temps = sem_df.groupby('sem_template').size().reset_index(name='count')
count_temps = count_temps.sort_values('count', ascending=False).reset_index(drop=True)
count_temps.head(10)
count_temps['frequency'] = count_temps['count'] / count_temps['count'].sum()
count_temps.head(10)
Explanation: Frequency of templates (identified by sentence number)
End of explanation
sem_df.head()
sem_df['sentence_number'].shape
temp_by_pos = pd.crosstab(sem_df['sem_template'], sem_df['sentence_number']).apply(lambda x: x / x.sum(), axis=0)
temp_by_pos.head()
Explanation: Frequency of templates by sentence position
End of explanation
df_dir = 'data/processed_dfs'
# Write sentences, mentions, predicates, and umls concepts to parquet, sem_df
sents_with_mentions.to_parquet(f'{df_dir}/sentences.parquet')
mentions.to_parquet(f'{df_dir}/mentions.parquet')
preds.to_parquet(f'{df_dir}/predicates.parquet')
umls.to_parquet(f'{df_dir}/umls.parquet')
sem_df.to_parquet(f'{df_dir}/templates.parquet')
temp_by_pos.to_parquet(f'{df_dir}/templates_by_pos.parquet')
Explanation: Write dataframes to parquet
We want to write these to a parquet file so that they can be used by a separate notebook to do clustering and note generation. This is just prep-work for those processes.
End of explanation
<END_TASK>