
    Pi                         d dl Z d dlmZmZ d dlmZmZmZ d dlm	Z	 d dl
Z
d dlZddlT  G d de          Z G d	 d
e          ZeZ G d de          ZdS )    N)ABCabstractmethod)OptionalSequenceTuple)OrderedDict   )*c                      e Zd ZdZddefdZeedefd                        Zde	edf         de
e	edf                  fd	Zedee         dd
fd            Zedee         defd            Zedee         dd
ddfd            ZdS )BaseLlamaCachez'Base cache class for a llama.cpp model.        capacity_bytesc                     || _         d S N)r   )selfr   s     i/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/llama_cpp/llama_cache.py__init__zBaseLlamaCache.__init__   s    ,    returnc                     t           r   NotImplementedErrorr   s    r   
cache_sizezBaseLlamaCache.cache_size   s
     "!r   key.c                     d S r    r   r   s     r   _find_longest_prefix_keyz'BaseLlamaCache._find_longest_prefix_key   s	     	r   llama_cpp.llama.LlamaStatec                     t           r   r   r   s     r   __getitem__zBaseLlamaCache.__getitem__"       !!r   c                     t           r   r   r   s     r   __contains__zBaseLlamaCache.__contains__&   r#   r   valueNc                     t           r   r   r   r   r&   s      r   __setitem__zBaseLlamaCache.__setitem__*   s
     "!r   r   )__name__
__module____qualname____doc__intr   propertyr   r   r   r   r   r   r"   boolr%   r)   r   r   r   r   r      sG       11- -s - - - - "C " " " ^ X"38_ 
%S/	"    "x} "1M " " " ^" " "$ " " " ^" "C=")E"	" " " ^" " "r   r   c                        e Zd ZdZddef fdZed             Zdeedf         de	eedf                  fd	Z
dee         dd
fdZdee         defdZdee         dd
fdZ xZS )LlamaRAMCachez&Cache for a llama.cpp model using RAM.r   r   c                 ~    t                                          |           || _        t                      | _        d S r   )superr   r   r   cache_state)r   r   	__class__s     r   r   zLlamaRAMCache.__init__4   s:    (((, MM 	r   c                 b    t          d | j                                        D                       S )Nc                     g | ]	}|j         
S r   )llama_state_size).0states     r   
<listcomp>z,LlamaRAMCache.cache_size.<locals>.<listcomp>=   s    RRRuE*RRRr   )sumr6   valuesr   s    r   r   zLlamaRAMCache.cache_size;   s.    RR8H8O8O8Q8QRRRSSSr   r   .r   c                 ~    d}d }fd| j                                         D             }|D ]\  }}||k    r|}|}|S )Nr   c              3   f   K   | ]+}|t           j        j                            |          fV  ,d S r   )	llama_cppllamaLlamalongest_token_prefix)r;   kr   s     r   	<genexpr>z9LlamaRAMCache._find_longest_prefix_key.<locals>.<genexpr>E   sQ       
 
 	%::1cBBC
 
 
 
 
 
r   )r6   keys)r   r   min_lenmin_keyrH   rF   
prefix_lens    `     r   r   z&LlamaRAMCache._find_longest_prefix_key?   sv     
 
 
 
%**,,
 
 
 " 	 	MAzG##$r   r    c                     t          |          }|                     |          }|t          d          | j        |         }| j                            |           |S NzKey not found)tupler   KeyErrorr6   move_to_endr   r   _keyr&   s       r   r"   zLlamaRAMCache.__getitem__O   s\    Cjj,,S11<?+++ &$$T***r   c                 J    |                      t          |                    d uS r   r   rN   r   s     r   r%   zLlamaRAMCache.__contains__X   !    ,,U3ZZ88DDr   r&   c                 <   t          |          }|| j        v r| j        |= || j        |<   | j        | j        k    r]t	          | j                  dk    rI| j                            d           | j        | j        k    rt	          | j                  dk    Cd S d S d S d S )Nr   F)last)rN   r6   r   r   lenpopitemr(   s      r   r)   zLlamaRAMCache.__setitem__[   s    Cjj$""" % %o 333D<L8M8MPQ8Q8Q$$%$000 o 333D<L8M8MPQ8Q8Q8Q8Q33338Q8Qr   r*   )r+   r,   r-   r.   r/   r   r0   r   r   r   r   r   r"   r1   r%   r)   __classcell__r7   s   @r   r3   r3   1   s"       00 s       T T XT38_ 
%S/	"    x} 1M    E E$ E E E E1x} 15Q 1 1 1 1 1 1 1 1r   r3   c                        e Zd ZdZ	 ddedef fdZed             Zde	ed	f         d
e
e	ed	f                  fdZdee         d
dfdZdee         d
efdZdee         ddfdZ xZS )LlamaDiskCachez'Cache for a llama.cpp model using disk..cache/llama_cacher   	cache_dirr   c                 |    t                                          |           t          j        |          | _        d S r   )r5   r   	diskcacheCachecache)r   r_   r   r7   s      r   r   zLlamaDiskCache.__init__k   s2     	(((_Y//


r   c                 N    t          | j                                                  S r   )r/   rc   volumer   s    r   r   zLlamaDiskCache.cache_sizeq   s    4:$$&&'''r   r   .r   c                     d}d }| j                                         D ]1}t          j        j                            ||          }||k    r|}|}2|S )Nr   )rc   iterkeysrB   rC   rD   rE   )r   r   rI   rJ   rF   rK   s         r   r   z'LlamaDiskCache._find_longest_prefix_keyu   s`     -1$$&& 	 	A".CCAsKKJG##$r   r    c                     t          |          }|                     |          }|t          d          | j                            |          }|S rM   )rN   r   rO   rc   poprQ   s       r   r"   zLlamaDiskCache.__getitem__   sM    Cjj,,S11<?+++.2jnnT.B.B r   c                 J    |                      t          |                    d uS r   rT   r   s     r   r%   zLlamaDiskCache.__contains__   rU   r   r&   c                 $   t          dt          j                   t          |          }|| j        v r#t          dt          j                   | j        |= || j        |<   t          dt          j                   | j        | j        k    rit          | j                  dk    rQt          t          | j                            }| j        |= | j        | j        k    rt          | j                  dk    Qt          dt          j                   d S )Nz"LlamaDiskCache.__setitem__: called)filez"LlamaDiskCache.__setitem__: deletezLlamaDiskCache.__setitem__: setr   z LlamaDiskCache.__setitem__: trim)
printsysstderrrN   rc   r   r   rX   nextiter)r   r   r&   key_to_removes       r   r)   zLlamaDiskCache.__setitem__   s    2DDDDCjj$*6SZHHHH
3
3/cjAAAAo 333DJ!8K8K dj!1!122M
=) o 333DJ!8K8K 	0szBBBBBBr   )r^   r   )r+   r,   r-   r.   strr/   r   r0   r   r   r   r   r   r"   r1   r%   r)   rZ   r[   s   @r   r]   r]   h   s9       11 MT0 00EH0 0 0 0 0 0 ( ( X(38_ 
%S/	"   	x} 	1M 	 	 	 	E E$ E E E ECx} C5Q C C C C C C C Cr   r]   )rn   abcr   r   typingr   r   r   collectionsr   ra   llama_cpp.llamarB   llama_typesr   r3   
LlamaCacher]   r   r   r   <module>rz      s+   



 # # # # # # # #         
 $ # # # # #            " " " " "S " " "@01 01 01 01 01N 01 01 01h 
3C 3C 3C 3C 3C^ 3C 3C 3C 3C 3Cr   