
    dh                        S SK Jr  S SKrS SKJrJrJrJr  S SK	r
S SKJr  S SKJr  S SKJr  S SKJr  S SKJr  SS	 jr " S
 S\5      rg)    )annotationsN)AnyIterableListOptional)CallbackManagerForRetrieverRun)Document)
Embeddings)BaseRetriever)
ConfigDictc           	         [         R                  R                  5        n[        R                  " [        UR                  UR                  U 5      5      5      sSSS5        $ ! , (       d  f       g= f)z
Create an index of embeddings for a list of contexts.

Args:
    contexts: List of contexts to embed.
    embeddings: Embeddings model to use.

Returns:
    Index of embeddings.
N)
concurrentfuturesThreadPoolExecutornparraylistmapembed_query)contexts
embeddingsexecutors      ]/var/www/html/shao/venv/lib/python3.13/site-packages/langchain_community/retrievers/nanopq.pycreate_indexr      sF     
			.	.	0HxxX\\**@*@(KLM 
1	0	0s   9A""
A0c                     \ rS rSr% SrS\S'    SrS\S'    S\S	'    SrS
\S'    SrS\S'    Sr	S\S'    Sr
S\S'    SrS\S'    \" SS9r\ S         SS jj5       r\        SS j5       r      SS jrSrg)NanoPQRetriever   z`NanoPQ retriever.r
   r   Nr   index	List[str]textsOptional[List[dict]]	metadatas   intkzOptional[float]relevancy_thresholdsubspace   clustersT)arbitrary_types_allowedc                2    [        X5      nU " SUUUUS.UD6$ )N)r   r   r    r"    )r   )clsr    r   r"   kwargsr   s         r   
from_textsNanoPQRetriever.from_texts5   s6     U/ 
!	

 
 	
    c                N    [        S U 5       6 u  pEU R                  " SXBUS.UD6$ )Nc              3  N   #    U H  oR                   UR                  4v   M     g 7fNpage_contentmetadata).0ds     r   	<genexpr>1NanoPQRetriever.from_documents.<locals>.<genexpr>M   s      Qy!..!**!=ys   #%)r    r   r"   r,   )zipr/   )r-   	documentsr   r.   r    r"   s         r   from_documentsNanoPQRetriever.from_documentsF   s:      Qy QR~~ 
)
GM
 	
r1   c                   SSK Jn  [        R                  " U R
                  R                  U5      5      n U" U R                  U R                  SS9R                  U R                  R                  S5      5      nUR!                  U R                  R                  S5      S
9nUR#                  UR                  S5      S9nUR%                  US9n	[        R&                  " U	5      n
U
SU R(                    Vs/ sH:  n[+        U R,                  U   U R.                  (       a  U R.                  U   O0 S9PM<     nnU$ ! [         a    [        S5      ef = f! [         a`    SR                  U R                  R                  S   U R                  U R                  U R                  R                  S   S	9n[        U5      ef = fs  snf )Nr   )PQzBCould not import nanopq, please install with `pip install nanopq`.T)MKsverbosefloat32zReceived params: training_sample={training_sample}, n_cluster={n_clusters}, subspace={subspace}, embedding_shape={embedding_shape}. Issue with the combination. Please retrace back to find the exact error   )training_sample
n_clustersr'   embedding_shape)vecs)query)codesr5   )nanopqrA   ImportErrorr   r   r   r   r'   r)   fitr   astypeAssertionErrorformatshapeRuntimeErrorencodedtableadistargsortr%   r	   r    r"   )selfrK   run_managerrA   query_embedspqerror_message
index_codedtdists	sorted_ixrowtop_k_resultss                r   _get_relevant_documents'NanoPQRetriever._get_relevant_documentsR   s   	! xx ; ;E BC	.dmmtDHH

!!),B" YYDJJ$5$5i$@YA
YY\00;Y<z*JJu%	 !TVV,

 -	 !ZZ_04,B - 	 
 M  	T 	  	.> f $

 0 0 3== $

 0 0 3	    }--	.(
s   D< AE 8A G<EA*F?r,   r4   )
r    r   r   r
   r"   r!   r.   r   returnr   )r=   zIterable[Document]r   r
   r.   r   rf   r   )rK   strrZ   r   rf   zList[Document])__name__
__module____qualname____firstlineno____doc____annotations__r   r"   r%   r&   r'   r)   r   model_configclassmethodr/   r>   rd   __static_attributes__r,   r1   r   r   r      s   "E3!&*I#*9AsJ&+//"HcPHc& $L 
 +/	

 
 (	

 
 

 
  	
%	
 	
 		

 
	
 	
++*H+	+r1   r   )r   r   r   r
   rf   z
np.ndarray)
__future__r   concurrent.futuresr   typingr   r   r   r   numpyr   langchain_core.callbacksr   langchain_core.documentsr	   langchain_core.embeddingsr
   langchain_core.retrieversr   pydanticr   r   r   r,   r1   r   <module>rz      s7    "  0 0  C - 0 3 N`m `r1   