
    ug                        d dl mZmZ d dlZd dlZd dlZd dlmZm	Z	m
Z
mZmZ d dlmZ d dlmZ  G d de	e         Zdefd	Zdefd
Zej(                   e       fde	e   deej*                  ddf   fd       Z e       ddfdej*                  de	e   dededdf
dZej2                  j4                  dej*                  ddfd       Zy)    )	GeneratorcastN)
EmbeddableEmbeddingFunction
EmbeddingsImageDocument)hashing_embedding_function)_exact_distancesc                   $    e Zd ZddZdedefdZy)hashing_multimodal_efreturnNc                 D    t        dt        j                        | _        y )N
   )dimdtype)r
   npfloat64_hef)selfs    X/var/www/openai/venv/lib/python3.12/site-packages/chromadb/test/ef/test_multimodal_ef.py__init__zhashing_multimodal_ef.__init__   s    .2RZZH	    inputc                 
   |D cg c]  }t        |       }}t        j                  | j                  |            }|t        j                  j                  |dd      z  }t        t        |j                               S c c}w )N   T)axiskeepdims)	strr   arrayr   linalgnormr   r   tolist)r   r   ito_texts
embeddingss        r   __call__zhashing_multimodal_ef.__call__   sk    $)*EqCFE*XXdii12
 	biinnZa$nGG
J
 1 1 344 +s   B )r   N)__name__
__module____qualname__r   r   r   r'    r   r   r   r      s    I5j 5Z 5r   r   r   c                  d    t         j                  j                  dddt         j                        S )Nr      )r   r      )sizer   )r   randomrandintint64r+   r   r   random_imager3      s$    99Q+RXXFFr   c                  (    t        t                     S )N)r   r3   r+   r   r   random_documentr5   #   s    |~r   
default_efc              #      K   t        j                         }|j                  d|       }| |j                          y w)Nmultimodal_collection)nameembedding_function)chromadbClientcreate_collectionclear_system_cache)r6   client
collections      r   r8   r8   '   sB      __F))$ * J 
s   =?r   r.   r8   
n_examplesn_query_resultsc                 n   t         j                  j                         }t         j                  j                  d       t	        |      D cg c]  }t        |       }}t	        |      D cg c]  }t                }} ||      }	t	        |d|z        D cg c]  }t        |       }
}t	        |      D cg c]  }t                }} ||      }t        j                  t        d      5  | j                  |d   |d   |d          d d d        | j                  |
|       | j                  ||       | j                  dg	      }t        |d
         t        |
      t        |      z   k(  sJ t        |d
         D ]H  \  }}||
v s||v sJ |d   J ||
v r|d   |   ||
j                  |         k(  sJ ||v s?|d   |   HJ  t               } ||g      }t!        ||	|z         \  }}|d   d | D cg c]  }||k  r||   n|
||z      }}t               } ||g      }t!        ||	|z         \  }}|d   d | D cg c]  }||k  r||   n|
||z      }}t        j                  t              5  | j#                  |g|g       d d d        | j#                  |g|dg      }|d
   d   |k(  sJ | j#                  |g|dg      }|d
   d   |k(  sJ t         j                  j%                  |       y c c}w c c}w c c}w c c}w # 1 sw Y   	xY wc c}w c c}w # 1 sw Y   xY w)Nr      z|Exactly one of (images|documents|uris)(?:, (images|documents|uris))?(?:, (images|documents|uris))? must be provided in add\.)match)ids	documentsimagesrF   rG   rF   rH   rG   )includerF   )query_imagesquery_texts)rL   	n_resultsrK   )rM   rN   rK   )r   r0   	get_stateseedranger   r3   r5   pytestraises
ValueErroraddgetlen	enumerateindexr   query	set_state)r8   r6   rA   rB   random_stater$   	image_ids_rH   image_embeddingsdocument_idsrG   document_embeddings
get_resultidquery_imagequery_image_embeddingimage_neighbor_indicesnearest_image_neighbor_idsquery_documentquery_document_embeddingdocument_neighbor_indicesnearest_document_neighbor_idsquery_results                           r   test_multimodalrm   4   s    99&&(LIINN1!&z!23!2AQ!2I3&+J&78&7ln&7F8!&)$)*a*n$EF$EqCF$ELF,1*,=>,=q",=I>$Y/ 
 N

 	!!!	!VAY 	" 	

 ,)D)F; '**K=*AJz% !S%6Y%GGGG:e,-2\!R9_44+&222k*1-<;M;Mb;Q1RRRR?k*1-555 . .K&}5 0/2EE!A (*+;O<"<A J	!LZ,HH<  " %&N)>*:;#3 "25H"H$ q
 +1-.>?%?A J	!LZ,HH? " % 
z	"##%^4D 	$ 	
 
# )..!]o} / L q!%???? )..#$+ / L q!%BBBBII%] 48 G>
 
D"% 
#	"s<   L ,LL
9L/L>L!L&<L+L+L4c                     t               }t               }d}| j                  ||       | j                  ||       | j	                  |dg      }|d   J |d   d   J y )N0rI   rJ   rG   )rF   rK   r   )r5   r3   rU   updaterV   )r8   documentimagerc   rb   s        r   !test_multimodal_update_with_imagers      s|      HNE	B"9  R 6&**rK=*IJk"...k"1%---r   )typingr   r   numpyr   rR   r;   chromadb.api.typesr   r   r   r   r	   !chromadb.test.property.strategiesr
   !chromadb.test.property.invariantsr   r   r3   r5   fixture
Collectionr8   intrm   markxfailrs   r+   r   r   <module>r~      s   "     I >

5-j9 
5Ge G  0E0G !*- x""D$./    1F0G	X&#..X&!*-X& X& 	X&
 
X&v .#...	. .r   