alt_llama_cpp.py

#Qwen2.5-Coder-32B製。llama.cppのllama-serverをllama-cpp-pythonライクに操作するライブラリ
from openai import OpenAI
import subprocess
from typing import List, Optional

llama_server_path = "llama_cpp"#llama-server 実行ファイルのパスを指定
llama_server_process = None  # グローバル変数としてプロセスオブジェクトを保持

def start_llama_server(model_path: str, 
                       n_gpu_layers: Optional[int] = None, 
                       n_ctx: Optional[int] = 512, 
                       n_batch: Optional[int] = 512, 
                       flash_attn: Optional[bool] = None, 
                       no_mmap: Optional[bool] = None,
                       tensor_split: Optional[List[int]] = None, 
                       last_n_tokens_size: Optional[int] = 64,
                       override_tensor: Optional[str]= None) -> None:
    """
    llama-server.exe を起動する関数

    Args:
        model_path (str): モデルのファイルパス
        n_gpu_layers (Optional[int]): GPUレイヤーの数 (デフォルト: None)
        n_ctx (Optional[int]): コンテキストサイズ (デフォルト: 512)
        n_batch (Optional[int]): バッチサイズ (デフォルト: 512)
        flash_attn (Optional[bool]): Flash Attention を使用するかどうか (デフォルト: None)
        no_mmap (Optional[bool]):メモリマップを無効化します (デフォルト: None)
        tensor_split (Optional[List[int]]): テンソルの分割比率 (デフォルト: None)
        last_n_tokens_size (Optional[int]): 最後のnトークンのサイズ (デフォルト: 64)
        override_tensor (Optional[str]):テンソルバッファタイプをオーバーライドします (デフォルト: None)
    """
    global llama_server_process  # グローバル変数を使用
    
    command = [
        f"{llama_server_path}\\llama-server.exe",
        "-m", model_path
    ]
    
    if n_gpu_layers is not None:
        if n_gpu_layers == -1:
            command.extend(["-ngl", "9999"])
        else:
            command.extend(["-ngl", str(n_gpu_layers)])
    
    if n_ctx is not None:
        command.extend(["-c", str(n_ctx)])
    
    if n_batch is not None:
        command.extend(["-b", str(n_batch)])
    
    if flash_attn is True:
        command.append("-fa")

    if no_mmap is True:
        command.append("--no-mmap")

    if tensor_split is not None:
        command.extend(["-ts", ",".join(map(str, tensor_split))])
    
    if last_n_tokens_size is not None:
        command.extend(["--repeat-last-n", str(last_n_tokens_size)])

    if override_tensor is not None:
        command.extend(["-ot", str(override_tensor)])    
    try:
        llama_server_process = subprocess.Popen(
            command, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.STDOUT, 
            text=True,  # テキストモードで出力を扱う
            encoding='utf-8'  # UTF-8 でエンコード
        )
        # 出力をキャプチャして表示
        for line in llama_server_process.stdout:
            print(f"{line.strip()}")
            if line.strip() == "srv  update_slots: all slots are idle":
                #print("llama-server を正常に起動しました。")
                break
        return "load_complate"
    except Exception as e:
        print(f"llama-server の起動中にエラーが発生しました: {e}")
# 使用例
#start_llama_server(
#    model_path="F:\\gguf\\sarashina2.1-1b-sft-imatrix-Q8_0-4329.gguf",
#    n_gpu_layers=32,
#    n_ctx=512,
#    n_batch=512,
#    flash_attn=True,
#    no_mmap=True,
#    override_tensor="blk\\\.([0-9]*[05])\\\.ffn_.*_exps\\\.=CPU",
#    tensor_split=[100, 0],
#    last_n_tokens_size=64
#)

def stop_llama_server():
    """
    起動中の llama-server.exe を終了する関数
    """
    global llama_server_process  # グローバル変数を使用
    
    if llama_server_process is not None:
        llama_server_process.terminate()  # プロセスを終了
        llama_server_process = None  # プロセスオブジェクトをリセット
        #print("llama-server を正常に終了しました。")
    else:
        #print("llama-server は起動していません。")
        pass
# 使用例
#stop_llama_server()

def response(prompt,max_tokens=16,temperature=0.8, top_p=0.95, frequency_penalty=0.0, presence_penalty=0.0,seed=-1,stop=None, stream=False):
    """
    起動中の llama-server.exe に回答してもらう関数

    Args:
        prompt (str): プロント (デフォルト: None)
        max_tokens (int): 最大出力トークン (デフォルト: 16)
        temperature (Optional[float]): 温度 (デフォルト: 0.8)
        top_p (Optional[float]):  (デフォルト: 0.95)
        frequency_penalty (Optional[float]):  (デフォルト: 0.0)
        presence_penalty (Optional[float]):  (デフォルト: 0.0)
        seed (Optional[int]): シード値 (デフォルト: -1)
        stop (Optional[list[str]]): ストップトークン (デフォルト: None)
        stream (Optional[bool]): ストリーミング出力する (デフォルト: False)
    """
    client = OpenAI(
        api_key="YOUR_API_KEY",
        base_url="http://localhost:8080/v1"
    )
    
    messages = [
        {"role": "user", "content": prompt},
    ]
    
    stream_options = {"include_usage": True} if stream else None
    
    response = client.chat.completions.create(
        model="gpt-4",
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        frequency_penalty=frequency_penalty,
        presence_penalty=presence_penalty,
        seed=seed,
        stop=stop if stop is not None else [],
        stream=stream,
        stream_options=stream_options,
    )
    
    return response

#使用例1(通常出力)
#system_prompt="あなたは優秀で誠実なアシスタントです"
#prompt = system_prompt+"こんにちは"
#response = alt_llama_cpp.response(prompt,
#                                  stream=False)
#通常出力の処理
#output = response.choices[0].message.content
#print(output)
#
#alt_llama_cpp.stop_llama_server()#llama-serverを終了してVRAM/RAMを開放する
#使用例2(ストリーミング出力)
#system_prompt="あなたは優秀で誠実なアシスタントです"
#prompt = system_prompt+"東京の魅力を100文字で教えて"
#response = alt_llama_cpp.response(prompt,
#                                  max_tokens=512,
#                                  temperature = 0.6,
#                                  top_p=0.95, 
#                                  stop=["<|stop|>"] ,
#                                  stream=True)
#ストリーミング出力の処理
#for chunk in response:
#    output =  chunk.choices[0].delta.content
#    if output:
#       print(output,end="",flush=True)
#alt_llama_cpp.stop_llama_server()#llama-serverを終了してVRAM/RAMを開放する