
    Ig              	           d dl mZmZmZmZ d dlZd dlmZ d dl	m
Z
 ddddd	d
diddgdgZ G d de      Z G d de
      Zy)    )AnyDictListTypeN)Colbert)TextEmbeddingWorkerzjinaai/jina-colbert-v2   ziNew model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 yearzcc-by-nc-4.0gQ@hfzonnx/model.onnxzonnx/model.onnx_data)modeldimdescriptionlicense
size_in_GBsources
model_fileadditional_filesc                        e Zd ZdZdZdZdZedee	   fd       Z
edeeeef      fd       Z	 ddeeej"                  f   d	ed
edeeej"                  f   f fdZ xZS )JinaColberti i    z<mask>returnc                     t         S )N)JinaColbertEmbeddingWorkerclss    d/var/www/html/answerous/venv/lib/python3.12/site-packages/fastembed/late_interaction/jina_colbert.py_get_worker_classzJinaColbert._get_worker_class   s    ))    c                     t         S )zLists the supported models.

        Returns:
            List[Dict[str, Any]]: A list of dictionaries containing the model information.
        )supported_jina_colbert_modelsr   s    r   list_supported_modelsz!JinaColbert.list_supported_models#   s
     -,r   
onnx_inputis_dockwargsc                 <    t         |   ||      }|sd|d   d d  |S )N   attention_mask)super_preprocess_onnx_input)selfr!   r"   r#   	__class__s       r   r(   z"JinaColbert._preprocess_onnx_input,   s1     W3JG
 ./J'(+r   )T)__name__
__module____qualname__QUERY_MARKER_TOKEN_IDDOCUMENT_MARKER_TOKEN_IDMIN_QUERY_LENGTH
MASK_TOKENclassmethodr   r   r   r   r   strr   r    npndarrayboolr(   __classcell__)r*   s   @r   r   r      s    "%J*$':"; * * -d4S>&: - - AEsBJJ/9=PS	c2::o	 r   r   c                        e Zd ZdededefdZy)r   
model_name	cache_dirr   c                      t        d||dd|S )Nr%   )r9   r:   threads )r   )r)   r9   r:   r#   s       r   init_embeddingz)JinaColbertEmbeddingWorker.init_embedding8   s'     
!
 	
 	
r   N)r+   r,   r-   r3   r   r>   r=   r   r   r   r   7   s    
 
 
; 
r   r   )typingr   r   r   r   numpyr4   "fastembed.late_interaction.colbertr   fastembed.text.onnx_text_modelr   r   r   r   r=   r   r   <module>rC      sd    ( (  6 >
 * C!*
 (34!  ' <
!4 
r   