hymenjj/llama-cpp-python-prebuilt
0
1�
2 00�i( � �p � d dl Z d dlmZ d dlZd dlmZ G d� de j � � Z G d� de� � Z dS )� N)�Anyc �| � e Zd Zej dej ej de dej ej fd�� � Z3dS )�LlamaDraftModel� input_ids�kwargs�returnc � � t � � ��N)�NotImplementedError��selfr r s �F/usr/local/lib/python3.11/site-packages/llama_cpp/llama_speculative.py�__call__zLlamaDraftModel.__call__4 s � � "�#�#�#� N) �__name__�5__module__� __qualname__�abc�abstractmethod�npt�NDArray�np�intcr r � r r r r s` � � � � � ���$� ��R�W�-�$�<?�$� ��R�W� �$� $� $� ��$� $� $r r c � � e Zd ZdZd
dedefd�Zedej e j6 dedefd�� � Z dej e j7 d e d8ej e j9 fd �Z
d S )�LlamaPromptLookupDecodingz>Based on https://github.com/apoorvumang/prompt-lookup-decoding� �10 �max_ngram_size�num_pred_tokensc �" � || _ || _ d S r11 )r r )r
r r s r �__init__z"LlamaPromptLookupDecoding.__init__ s � � ,��� .�� � � r r c �� � | j d }t t ||dz12 � � dd� � D ]�}t j j � | |f� � }| | d � }t j ||k d�� � }t j |� � d }|D ]0} | |z }13|14|z } t | |� � } |15| k r| |16| � c c S �1��t j g t j17 �� � S )Nr � �����)�axis)�dtype) �shape�range�minr �lib�
stride_tricks�sliding_window_view�all�nonzero�arrayr ) r r r � input_length�18ngram_size�windows� ngram_array�matches�
match_indices�idx� start_idx�end_idxs r �find_candidate_pred_tokensz4LlamaPromptLookupDecoding.find_candidate_pred_tokens s � � !��q�)� � ��N�L�1�4D� E� E�q�"�M�M� 8� 8�J��f�*�>�>�y�:�-�X�X�G� $�Z�K�L�L�1�K� �f�W� � 3�!�<�<�<�G� �J�w� /� /�� 2�M� %�
8�
8�� �*� ,� �#�o�5�� �g�|�4�4�� �w�&�&�$�Y�w�%6�7�7�7�7�7�7� '�
8� �x��"�'�*�*�*�*r r r c �F � | � || j | j �� � S )N)r r r )r: r r r s r r z"LlamaPromptLookupDecoding.__call__9 s0 � � �.�.� � �.� � 0� /� 19� 20� 21r N)r r )r r r �__doc__�intr" � staticmethodr r r r r: r r r r r r r s� � � � � � �H�H�/� /�s� /�� /� /� /� /� � +��;�r�w�'� +�� +� � +� +� +� �\� +�@22� ��R�W�-�23�<?�24� ��R�W� �25� 26� 27� 28� 29� 30r r )31r �typingr �numpyr � numpy.typingr �ABCr r r r r �<module>rC s� �� 32�33�34�35� � � � � � � � � � � � � � � � �$� $� $� $� $�c�g� $� $� $�/36� /37� /38� /39� /40�� /41� /42� /43� /44� /45r