Script to build a version of llama-cpp-python for AMD renamed in accordance with Oobabooga's odd requirements

build_renamed_rocm_wheel.py

"""
build_renamed_rocm_wheel.py
===========================

:author: David Griffin <[email protected]>
:license: MIT License
:copyright: 2024


Copyright 2024, David Griffin

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

v241010: Replaced Path.walk with os.walk for compatibility with Python 3.11
"""

import pathlib
import subprocess
import sys
import os

f_output_path = pathlib.Path('llama-cpp-python-cuda')
if not f_output_path.exists():
    f_output_path.mkdir()
base_path = pathlib.Path('llama-cpp-python')
renames = {'llama_cpp': 'llama_cpp_cuda', 'llama-cpp-python': 'llama-cpp-python-cuda', 'llama_cpp.py': 'llama_cpp_cuda.py'}
ignore_dirs = {'.git', '.github'}

if not base_path.exists():
    print('Initial clone of llama-cpp-python')
    r = subprocess.run(['git', 'clone', '--recurse-submodules', 'https://github.com/abetlen/llama-cpp-python.git'], check=True)
else:
    print('Updating llama-cpp-python and dependencies')
    subprocess.run(['git', 'pull', '--recurse-submodules'], cwd=base_path, check=True)

print('Rewriting llama-cpp-python into llama-cpp-python-cuda')

def rewrite_dir(dirpath):
    outpath = []
    for part in dirpath.parts:
        outpath.append(renames.get(part, part))
    return pathlib.Path(*outpath)
    
for dirpath, dirnames, filenames in os.walk(base_path):
    dirpath = pathlib.Path(dirpath)
    for ignore in ignore_dirs:
        if ignore in dirnames: dirnames.remove(ignore)
    #print(repr(dirpath), dirnames, filenames)
    output_path = rewrite_dir(dirpath)
    if not output_path.exists():
        output_path.mkdir()
    for fn in filenames:
        with open(dirpath / fn, 'rb') as f:
            contents = f.read()
        if fn.endswith('.py') or fn.endswith('.toml') or fn.endswith('.txt'):
            contents = contents.replace(b'llama_cpp', b'llama_cpp_cuda')
            contents = contents.replace(b'llama_cpp_cuda_python', b'llama_cpp_python_cuda')
        
        fn = renames.get(fn, fn)
        with open(output_path / fn, 'wb') as f:
            f.write(contents)
            
with open(f_output_path / 'llama_cpp_cuda/llama_cpp.py', 'w') as f:
    f.write('from .llama_cpp_cuda import *')

print('Building ROCM Wheel')
build_env = os.environ.copy()
build_env['CMAKE_ARGS'] = '-DGGML_HIPBLAS=on'
subprocess.run([sys.executable, '-m', 'pip', 'wheel', '.'], env=build_env, cwd=f_output_path)