U
    Vc$                     @   s   d Z ddlZddlmZ ddlmZ ddlmZ ddlm	Z	 ddl
mZ ddlmZ d	Zd
d Zdd Zdd ZG dd dejZdS )zThe nPMI visualization plugin.    N)wrappers)plugin_util)base_plugin)	http_util)provider)metadata   c                 C   s   t j| d|idddS )Nerrorapplication/jsoni  )code)r   Respond)requesterror_message r   H/tmp/pip-unpacked-wheel-g8kmtpbc/tensorboard/plugins/npmi/npmi_plugin.py_error_response   s    r   c                 C   s
   t | dS )Nzrun parameter is not provided)r   )r   r   r   r   _missing_run_error_response(   s    r   c                 C   s   dd | D S )Nc                 S   s0   g | ](}t |trt|nt|r(d n|qS )N)
isinstancelistconvert_nan_nonemathisnan).0er   r   r   
<listcomp>.   s   
z$convert_nan_none.<locals>.<listcomp>r   )Zarrr   r   r   r   -   s    r   c                       s   e Zd ZdZejZ fddZdd Zdd Z	dd	 Z
d
d Zdd Zdd Zdd Zdd Zejjdd Zejjdd Zejjdd Zejjdd Zejjdd Z  ZS )
NpmiPluginznPMI Plugin for Tensorboard.c                    sJ   t t| | |j| _|jp i | jt| _	|j
| _tjddd| _dS )zInstantiates the nPMI Plugin via Tensorboard core.

        Args:
            context: A base_plugin.TBContext instance.
        ZnPMIr   )Z	data_kindZlatest_known_versionN)superr   __init__ZlogdirZ_logdirZsampling_hintsgetplugin_name_DEFAULT_DOWNSAMPLING_downsample_toZdata_provider_data_providerr   Z_MetadataVersionChecker_version_checker)selfcontext	__class__r   r   r   =   s    
 zNpmiPlugin.__init__c                 C   s   | j | j| j| j| jdS )N)z/tagsz/annotationsz/metricsz/valuesz/embeddings)
serve_tagsserve_annotationsserve_metricsserve_valuesserve_embeddingsr$   r   r   r   get_plugin_appsN   s    zNpmiPlugin.get_plugin_appsc                 C   s   dS )zDetermines whether this plugin is active.

        This plugin is only active if TensorBoard sampled any npmi summaries.

        Returns:
          Whether this plugin is active.
        Fr   r-   r   r   r   	is_activeW   s    zNpmiPlugin.is_activec                 C   s   t jddddS )NTZnpmi)Zis_ng_componentZtab_nameZdisable_reload)r   ZFrontendMetadatar-   r   r   r   frontend_metadataa   s
      zNpmiPlugin.frontend_metadatac                 C   s   | j j||| jd}dd |D }| D ]\\}}g ||< | D ]B\}}t|j}	| j|	j	||sjqBt|j}
|| 
| qBq*|S )N)experiment_idr   c                 S   s   i | ]
}|i qS r   r   r   runr   r   r   
<dictcomp>j   s      z(NpmiPlugin.tags_impl.<locals>.<dictcomp>)r"   list_tensorsr   itemsr   Zparse_plugin_metadataZplugin_contentr#   okversionappend)r$   ctx
experimentmappingresultr3   Ztag_to_contenttagZ	metadatumZmdcontentr   r   r   	tags_implf   s"      zNpmiPlugin.tags_implc           
   	   C   s   | j j||| jtjtjgdd}dd |D }| D ]d\}}| j j||| jtj|gtjgd| j	d}|
|i 
tji }dd |d	 jD }	|	||< q8|S )
Ntagsr1   r   run_tag_filterc                 S   s   i | ]
}|i qS r   r   r2   r   r   r   r4      s      z/NpmiPlugin.annotations_impl.<locals>.<dictcomp>runsrB   r1   r   rD   Z
downsamplec                 S   s   g | ]}| d qS zutf-8decode)r   
annotationr   r   r   r      s   z/NpmiPlugin.annotations_impl.<locals>.<listcomp>r   )r"   r5   r   r   RunTagFilterr   ZANNOTATIONS_TAGr6   read_tensorsr!   r   numpy)
r$   r:   r;   r<   r=   r3   _Zall_annotationsannotations
event_datar   r   r   annotations_implw   s:     	 
zNpmiPlugin.annotations_implc           
   	   C   s   | j j||| jtjtjgdd}dd |D }| D ]d\}}| j j||| jtj|gtjgd| j	d}|
|i 
tji }dd |d	 jD }	|	||< q8|S )
NrA   rC   c                 S   s   i | ]
}|i qS r   r   r2   r   r   r   r4      s      z+NpmiPlugin.metrics_impl.<locals>.<dictcomp>rE   rG   c                 S   s   g | ]}| d qS rH   rI   )r   Zmetricr   r   r   r      s     z+NpmiPlugin.metrics_impl.<locals>.<listcomp>r   )r"   r5   r   r   rL   r   ZMETRICS_TAGr6   rM   r!   r   rN   )
r$   r:   r;   r<   r=   r3   rO   Zall_metricsZmetricsrQ   r   r   r   metrics_impl   s,     	
zNpmiPlugin.metrics_implc           
   	   C   s   | j j||| jtjtjgdd}dd |D }| D ]f\}}| j j||| jtj|gtjgd| j	d}|
|i 
tji }|d j }	t|	}	|	||< q8|S )NrA   rC   c                 S   s   i | ]
}|i qS r   r   r2   r   r   r   r4      s      z*NpmiPlugin.values_impl.<locals>.<dictcomp>rE   rG   r   )r"   r5   r   r   rL   r   Z
VALUES_TAGr6   rM   r!   r   rN   tolistr   )
r$   r:   r;   r<   r=   r3   rO   Z
all_valuesvaluesrQ   r   r   r   values_impl   s.     	
zNpmiPlugin.values_implc           
   	   C   s   | j j||| jtjtjgdd}dd |D }| D ]^\}}| j j||| jtj|gtjgd| j	d}|
|i 
tji }|d j }	|	||< q8|S )NrA   rC   c                 S   s   i | ]
}|i qS r   r   r2   r   r   r   r4      s      z.NpmiPlugin.embeddings_impl.<locals>.<dictcomp>rE   rG   r   )r"   r5   r   r   rL   r   ZEMBEDDINGS_TAGr6   rM   r!   r   rN   rT   )
r$   r:   r;   r<   r=   r3   rO   Zall_embeddingsZ
embeddingsrQ   r   r   r   embeddings_impl   s6     	 
zNpmiPlugin.embeddings_implc                 C   s4   t |j}t |j}| j||d}t||dS N)r;   r
   )r   r%   environr1   r@   r   r   r$   r   r:   r;   contentsr   r   r   r(      s    zNpmiPlugin.serve_tagsc                 C   s4   t |j}t |j}| j||d}t||dS rX   )r   r%   rY   r1   rR   r   r   rZ   r   r   r   r)      s    zNpmiPlugin.serve_annotationsc                 C   s4   t |j}t |j}| j||d}t||dS rX   )r   r%   rY   r1   rS   r   r   rZ   r   r   r   r*      s    zNpmiPlugin.serve_metricsc                 C   s4   t |j}t |j}| j||d}t||dS rX   )r   r%   rY   r1   rV   r   r   rZ   r   r   r   r+      s    zNpmiPlugin.serve_valuesc                 C   s4   t |j}t |j}| j||d}t||dS rX   )r   r%   rY   r1   rW   r   r   rZ   r   r   r   r,      s    zNpmiPlugin.serve_embeddings)__name__
__module____qualname____doc__r   ZPLUGIN_NAMEr   r   r.   r/   r0   r@   rR   rS   rV   rW   r   RequestZapplicationr(   r)   r*   r+   r,   __classcell__r   r   r&   r   r   8   s*   	




r   )r_   r   Zwerkzeugr   Ztensorboardr   Ztensorboard.pluginsr   Ztensorboard.backendr   Ztensorboard.datar   Ztensorboard.plugins.npmir   r    r   r   r   ZTBPluginr   r   r   r   r   <module>   s   	