Run language models on consumer hardware.
from .download import download
from .llm import llm
from .uninstall import uninstall
from .install import install
ooba/__init__.py
KillianLucas/ooba
from .utils.get_app_dir import get_app_dir
from .utils.ensure_repo_exists import ensure_repo_exists
from .utils.list_gguf_files import list_gguf_files
from huggingface_hub import hf_hub_download
import os

DEFAULT_MODEL_DIR = os.path.join(get_app_dir(), "models")

def download(url, path=None, gguf_quality=0.4):

    if 'gguf' in url.lower():
        if path == None:
            path = DEFAULT_MODEL_DIR

        final_path = download_gguf(url, path, gguf_quality)
    else:
        final_path = download_non_gguf(url, path)

    return final_path

def download_gguf(url, path, gguf_quality):
    url_parts = url.split('/')
    repo_id = "/".join(url_parts[-2:])
    
    gguf_files = list_gguf_files(repo_id)
    index = int(gguf_quality * len(gguf_files))
    model_name = gguf_files[index]["filename"]

    final_path = os.path.join(path, model_name)

    if os.path.exists(final_path):
        print(f"Located '{final_path}', which has already been downloaded.")
        return final_path

    hf_hub_download(
        repo_id=repo_id,
        filename=model_name,
        local_dir=path,
        local_dir_use_symlinks=False,
        resume_download=True
    )

    return final_path

def download_non_gguf(url, path):
    ensure_repo_exists()

    script_path = os.path.join(get_app_dir(), 'text-generation-ui', 'download-model.py')
    os.system(f'python {script_path} {url} --output "{path}"')

    model_name = url.split("/")[-1]
    final_path = os.path.join(path, model_name)

    return final_path
ooba/download.py
KillianLucas/ooba
import os
import subprocess
from tqdm import tqdm
from .utils.detect_hardware import detect_hardware
from .utils.ensure_repo_exists import ensure_repo_exists
from .utils.get_app_dir import get_app_dir
from .uninstall import uninstall

REPO_DIR = os.path.join(get_app_dir(), 'text-generation-ui')

run_cmd_wrapper = """def run_cmd(cmd, **kwargs):
    cmd = cmd.replace("python ", '"' + sys.executable + '" ')
    print("Running command:", cmd)
    return wrapped_run_cmd(cmd, **kwargs)

def wrapped_run_cmd"""

def install(force_reinstall=False, verbose=False, cpu=False, gpu_choice=None, full_auto=False):
    """
    full_auto = if it can't find GPU, it will default to CPU. if this is false, we'll ask the user what to do if we can't detect a GPU
    """

    if not force_reinstall:
        installer_files_path = os.path.join(REPO_DIR, 'installer_files')
        if os.path.exists(installer_files_path):
            if verbose:
                print("Oobabooga is already installed.")
            return
    else:
        uninstall(confirm=False, entire_repo=True)

    detected_gpu_choice, start_script_filename = detect_hardware()
    start_script = os.path.join(REPO_DIR, start_script_filename)

    if cpu:
        gpu_choice = "N"

    if not full_auto:
        # If we couldn't find a GPU and the user didn't set one,
        if detected_gpu_choice == "N" and not gpu_choice:

            # Prompt user for GPU choice if no supported GPU is detected

            print("Could not automatically detect GPU. Please choose your GPU:\n")

            print("A) NVIDIA")
            print("B) AMD (Linux/MacOS only. Requires ROCm SDK 5.4.2/5.4.3 on Linux)")
            print("C) Apple M Series")
            print("D) Intel Arc (IPEX)")
            print("N) None (I want to run models in CPU mode)\n")

            gpu_choice = input("> ").upper()

            while gpu_choice not in 'ABCDN':
                print("Invalid choice. Please try again.")
                gpu_choice = input("> ").upper()

    if not gpu_choice:
        gpu_choice = detected_gpu_choice

    if verbose:
        print(f"Installing LLM interface package...")
    else:
        # Incinerate all stout
        #self.original_stdout = sys.stdout
        #sys.stdout = open(os.devnull, 'w')
        pass

    ensure_repo_exists(verbose=verbose)

    # I think this is actually for pytorch version, not for GPU, so I've commented it out:
    """
    if self.cpu:
        with open(os.path.join(REPO_DIR, 'one_click.py'), 'r') as file:
            filedata = file.read()

        # Force it to install CPU-only requirements
        filedata = filedata.replace('elif is_cpu:', 'elif True:')

        # Write the file out again
        with open(os.path.join(REPO_DIR, 'one_click.py'), 'w') as file:
            file.write(filedata)
    """

    # This lets us run one_click from conda env, not regular python
    for root, _, files in os.walk(REPO_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    filedata = f.read()
                filedata = filedata.replace('\npython one_click.py', '\n"$INSTALL_ENV_DIR/bin/python" one_click.py')
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write(filedata)
            except:
                continue


    with open(os.path.join(REPO_DIR, 'one_click.py'), 'r') as file:
            filedata = file.read()
    # Force all commands in one_click to be run from conda's python
    if "wrapped_run_cmd" not in filedata:
        filedata = filedata.replace("def run_cmd", run_cmd_wrapper)
    # weird little thing that fixed it for me
    if "git checkout main" not in filedata:
        git_checkout_main = 'run_cmd("git checkout main", assert_success=False, environment=True)'
        git_pull_autostash = 'run_cmd("git pull --autostash", assert_success=False, environment=True)'
        existing_line = "run_cmd(\"git pull --autostash\", assert_success=True, environment=True)"
        lines = filedata.split('\n')
        for i, line in enumerate(lines):
            if existing_line in line:
                indentation = len(line) - len(line.lstrip())
                lines[i] = f"{indentation * ' '}{git_checkout_main}\n{indentation * ' '}{git_pull_autostash}"
        filedata = '\n'.join(lines)
    with open(os.path.join(REPO_DIR, 'one_click.py'), 'w') as file:
        file.write(filedata)


    #### Run install command

    base_cmd = [
        start_script
    ]

    if gpu_choice == "N":
        base_cmd += ["--cpu"] # I'm not sure if this is necessary

    env = os.environ.copy()
    env["GPU_CHOICE"] = gpu_choice
    env["LAUNCH_AFTER_INSTALL"] = "False"
    env["INSTALL_EXTENSIONS"] = "False"

    update_cmd = base_cmd + ["--update"]

    print("Setting up the language model...\n\nThis can take up to 15 minutes. The progress bar might appear to freeze (some steps take several minutes).\n")

    # Initialize tqdm progress bar
    total_lines = 1738
    pbar = tqdm(total=total_lines, ncols=100)

    for cmd in [base_cmd, update_cmd]:
        process = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, bufsize=1)

        # Read the output line by line
        for line in iter(process.stdout.readline, ''):            
            if verbose:
                print(line.strip())
            else:
                # Update the progress bar by one step for each line
                pbar.update(1)

        process.wait()

    # After all processes are done, fill the progress bar to 100%
    if not verbose:
        pbar.n = total_lines
        pbar.refresh()

    # Close the progress bar
    pbar.close()
        
    if verbose:
        print("Install complete.")
ooba/install.py
KillianLucas/ooba
import os
import time
import asyncio
import json
import websockets
import subprocess
from .utils.get_open_ports import get_open_ports
from .utils.detect_hardware import detect_hardware
from .uninstall import uninstall
import random
import sys
import threading
import asyncio
import json
import websockets
from queue import Queue
from.utils.openai_messages_converters import role_content_to_history
from .install import install as install_oobabooga
from .utils.get_app_dir import get_app_dir

REPO_DIR = os.path.join(get_app_dir(), 'text-generation-ui')

class llm:
    def __init__(self, path, cpu=False, verbose=False, first_time=True):

        if first_time:
            print("\nGetting started...")

        try:
            self.path = path
            self.cpu = cpu
            self.verbose = verbose

            self.gpu_choice, start_script_filename = detect_hardware()

            self.start_script = os.path.join(REPO_DIR, start_script_filename)

            if cpu:
                self.gpu_choice = "N"

            # This will exit if it's already installed
            install_oobabooga(gpu_choice=self.gpu_choice)

            # Start oobabooga server
            model_dir = "/".join(path.split("/")[:-1])
            model_name = path.split("/")[-1]

            # Find an open port
            while True:
                self.port = random.randint(2000, 10000)
                open_ports = get_open_ports(2000, 10000)
                if self.port not in open_ports:
                    break

            # Also find an open port for blocking -- it will error otherwise
            while True:
                unused_blocking_port = random.randint(2000, 10000)
                open_ports = get_open_ports(2000, 10000)
                if unused_blocking_port not in open_ports:
                    break

            cmd = [
                self.start_script,
                "--model-dir", model_dir,
                "--model", model_name,
                "--api-blocking-port", str(unused_blocking_port),
                "--api-streaming-port", str(self.port),
                "--extensions", "api",
                "--api"
            ]

            if self.gpu_choice == "N":
                cmd.append("--cpu")
            if self.verbose:
                cmd.append("--verbose")

            if self.verbose:
                print("Running command:", " ".join(cmd))

            env = os.environ.copy()
            env["GPU_CHOICE"] = self.gpu_choice
            env["LAUNCH_AFTER_INSTALL"] = "True"
            env["INSTALL_EXTENSIONS"] = "False"
            
            if self.verbose:
                self.process = subprocess.Popen(cmd, env=env, stdout=sys.stdout, stderr=sys.stderr)
            else:
                self.process = subprocess.Popen(cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

            self.uri = f'ws://localhost:{self.port}/api/v1/chat-stream'

            # Wait for it to be ready by checking the port
            num_attempts = 50
            for attempt in range(num_attempts):
                open_ports = get_open_ports(2000, 10000)
                if self.port in open_ports:
                    if self.verbose:
                        print("Server is ready.")
                    break
                else:
                    if self.verbose:
                        print(f"Server is not ready... ({attempt+1}/{num_attempts})")
                    time.sleep(1.5)
            else:
                raise Exception("Server took too long to start")

            
            # Warm up the server (idk if it needs this)
            """
            if self.verbose:
                print("Warming up the server...")
                
            for _ in self.chat([{"role": "user", "content": "Hi"}], max_tokens=1):
                break

            if self.verbose:
                print("Server is warmed up.")
            """
            

        except Exception as e:
            print(e)

            if self.gpu_choice == "N":
                raise
            else:
                print("Auto GPU installation was unsuccessful. Re-installing for CPU use.")

                install_oobabooga(force_reinstall=True, cpu=True)
                self.__init__(path, cpu=True, verbose=self.verbose)

        if first_time:
            # Hack to fix it not working multiple times in a row. Must change this soon
            self.first_time = True


    def chat(self, messages, max_tokens=None, temperature=0):

        # Hack to fix it not working multiple times in a row. Must change this soon
        if self.first_time:
            self.first_time = False
        else:
            self.process.terminate()
            self.__init__(self.path, cpu=self.cpu, verbose=self.verbose, first_time=False)
                
        if any([message["role"] == "system" for message in messages[1:]]):
            raise ValueError("Only the first message can have {'role': 'system'}.")
        
        if "system" in messages[0]["role"]:
            system_message = messages[0]["content"]
            messages = messages[1:]
        else:
            system_message = "You are a helpful AI assistant."

        if messages[-1]["role"] != "user":
            raise ValueError("The last message role must be 'user'")

        user_input = messages[-1]["content"]
        messages = messages[:-1]

        messages = role_content_to_history(messages)

        request = {
            'context': system_message,
            'user_input': user_input,
            'history': messages,
            'mode': 'chat', # Valid options: 'chat', 'chat-instruct', 'instruct'
            'temperature': temperature,
        }

        if max_tokens == None:
            request["auto_max_new_tokens"] = True
        else:
            request["max_new_tokens"] = max_tokens
            request["auto_max_new_tokens"] = False

        q = Queue()

        def run_async_code():
            async def async_chat():
                async with websockets.connect(self.uri, ping_interval=None, open_timeout=9999) as websocket:
                    current_length = 0
                    await websocket.send(json.dumps(request))

                    while True:
                        try:
                            incoming_data = await websocket.recv()
                        except websockets.exceptions.ConnectionClosedError as e:
                            # This happens if they cancel mid-stream
                            if self.verbose:
                                print(e)
                            return
                        incoming_data = json.loads(incoming_data)

                        match incoming_data['event']:
                            case 'text_stream':
                                new_history = incoming_data['history']
                                if not new_history['visible']:
                                    continue
                                token = new_history['visible'][-1][1][current_length:]
                                current_length += len(token)
                                q.put(token)
                            case 'stream_end':
                                q.put(None)  # signal the end
                                return

            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(async_chat())

        t = threading.Thread(target=run_async_code)
        t.start()

        while True:
            item = q.get()
            if item is None:  # signal for completion
                break
            yield item
ooba/llm.py
KillianLucas/ooba

import ooba

ooba.uninstall(entire_repo=True)

path = ooba.download("https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF")

llm = ooba.llm(path, verbose=False, cpu=True)

messages = [
    {"role": "system", "content": "Your name is Shoggoth."},
    {"role": "user", "content": "What's your name?"}
]

for token in llm.chat(messages):
    print(token)
    break

print("Second one:")
for token in llm.chat(messages):
    print(token)
    break
ooba/test.py
KillianLucas/ooba
import os
import ooba
from ooba.utils.get_app_dir import get_app_dir

def test_ooba():
    if True:
        if os.path.exists(get_app_dir()):
            ooba.uninstall(confirm=False)

    url = "https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GGUF"

    path = ooba.download(url, gguf_quality=0)
    print("Downloaded to:", path)

    assert os.path.exists(path), "File not downloaded"

    llm = ooba.llm(path, verbose=True)

    messages = [{"role": "user", "content": "Hi Mistral."}]

    for token in llm.chat(messages):
        print(token)

    assert token
ooba/tests/test_ooba.py
KillianLucas/ooba
import shutil
import os
from ooba.utils.get_app_dir import get_app_dir

def uninstall(confirm=True, entire_repo=False):
    if entire_repo:
        repo_dir = os.path.join(get_app_dir(), 'text-generation-ui')
    else:
        repo_dir = os.path.join(get_app_dir(), 'text-generation-ui', 'installer_files')
    if confirm:
        user_input = input(f"Are you sure you want to uninstall? This will delete everything in `{repo_dir}` (y/n): ")
        if user_input.lower() != 'y':
            print("Uninstallation cancelled.")
            return
    if os.path.exists(repo_dir):
        shutil.rmtree(repo_dir)
    else:
        print(f"The directory `{repo_dir}` does not exist.")
ooba/uninstall.py
KillianLucas/ooba
import subprocess
import platform

def detect_hardware():

    #### Detect system

    system = platform.system()

    if system == "Linux":
        start_script_filename = "start_linux.sh"
    elif system == "Windows":
        start_script_filename = "start_windows.bat"
    elif system == "Darwin":
        start_script_filename = "start_macos.sh"
    else:
        print(f"OS {system} is likely not supported. Assuming system is Linux.")
        start_script_filename = "start_linux.sh"

    #### Detect GPU

    gpu_choice = detect_gpu()

    return gpu_choice, start_script_filename

def detect_gpu():
    """
    This returns a GPU choice in the way Oobabooga likes:

    Possible outputs:
    - "A" for NVIDIA GPU
    - "B" for AMD GPU
    - "C" for Apple M Series GPU
    - "D" for Intel Arc GPU
    - "N" if no supported GPU is detected
    """

    os_type = platform.system().lower()

    # Check for NVIDIA GPU
    try:
        nvidia_output = subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT, text=True)
        if "NVIDIA-SMI" in nvidia_output:
            return "A"
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass  # nvidia-smi command not found

    # Check for AMD GPU (on Linux)
    if os_type == 'linux':
        try:
            amd_output = subprocess.check_output(["lshw", "-C", "display"], stderr=subprocess.STDOUT, text=True)
            if "AMD" in amd_output:
                return "B"
        except (subprocess.CalledProcessError, FileNotFoundError):
            pass  # lshw command not found

    # Check for Apple M Series GPU (on MacOS)
    if os_type == 'darwin':
        try:
            metal_output = subprocess.check_output(["system_profiler", "SPDisplaysDataType"], stderr=subprocess.STDOUT, text=True)
            if "Apple M" in metal_output:
                return "C"
        except (subprocess.CalledProcessError, FileNotFoundError):
            pass  # system_profiler command not found

    # there isn't a reliable method to detect Intel Arc GPUs yet. - ChatGPT

    # GPU not properly detected. Output CPU option

    return "N"

ooba/utils/detect_hardware.py
KillianLucas/ooba
import os
from .get_app_dir import get_app_dir
import subprocess
from .get_latest_release import get_latest_release

REPO_DIR = os.path.join(get_app_dir(), 'text-generation-ui')
REPO_URL = "https://github.com/oobabooga/text-generation-webui.git"
TAG = get_latest_release()
TAG = "v1.7"

def ensure_repo_exists(verbose=False):
    try:
        if not os.path.isdir(REPO_DIR):
            subprocess.check_call(['git', 'clone', REPO_URL, REPO_DIR], stdout=None if verbose else subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            subprocess.check_call(['git', '-C', REPO_DIR, 'checkout', TAG], stdout=None if verbose else subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            if verbose:
                print("Successfully cloned repo.")
        else:
            if verbose:
                print("Repository already exists, skipping clone.")
    except subprocess.CalledProcessError:
        raise Exception("Failed to clone repository.")
ooba/utils/ensure_repo_exists.py
KillianLucas/ooba
import os
import appdirs
import platform

def get_app_dir():
    if platform.system() == "Darwin":  # macOS
        repo_dir = os.path.join(os.path.expanduser("~"), ".ooba")
    else:
        repo_dir = appdirs.user_data_dir("ooba")
    return repo_dir
ooba/utils/get_app_dir.py
KillianLucas/ooba
import requests

def get_latest_release():
    response = requests.get("https://api.github.com/repos/oobabooga/text-generation-webui/releases/latest")
    if response.status_code == 200:
        return response.json()["tag_name"]
    else:
        raise Exception("Failed to fetch latest release from GitHub API.")
ooba/utils/get_latest_release.py
KillianLucas/ooba
import socket

def get_open_ports(start_port=2000, end_port=10000, host='localhost'):
    open_ports = []
    for port in range(start_port, end_port+1):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            socket.setdefaulttimeout(1)
            result = s.connect_ex((host, port))  # returns 0 if the connection was successful
            if result == 0:
                open_ports.append(port)
    return open_ports
ooba/utils/get_open_ports.py
KillianLucas/ooba
from huggingface_hub import list_files_info, login
from typing import Dict, List, Union

def list_gguf_files(repo_id: str) -> List[Dict[str, Union[str, float]]]:
    """
    Fetch all files from a given repository on Hugging Face Model Hub that contain 'gguf'.

    :param repo_id: Repository ID on Hugging Face Model Hub.
    :return: A list of dictionaries, each dictionary containing filename, size, and RAM usage of a model.
    """

    try:
      files_info = list_files_info(repo_id=repo_id)
    except Exception as e:
      if "authentication" in str(e).lower():
        print("You likely need to be logged in to HuggingFace to access this language model.")
        print(f"Visit this URL to log in and apply for access to this language model: https://huggingface.co/{repo_id}")
        print("Then, log in here:")
        login()
        files_info = list_files_info(repo_id=repo_id)

    gguf_files = [file for file in files_info if "gguf" in file.rfilename]

    gguf_files = sorted(gguf_files, key=lambda x: x.size)

    # Prepare the result
    result = []
    for file in gguf_files:
        size_in_gb = file.size / (1024**3)
        filename = file.rfilename
        result.append({
            "filename": filename,
            "Size": size_in_gb,
            "RAM": size_in_gb + 2.5,
        })

    return result
ooba/utils/list_gguf_files.py
KillianLucas/ooba
def history_to_role_content(history):
    role_content_list = []
    for visible_interaction in history['visible']:
        user_content, assistant_content = visible_interaction
        role_content_list.append({'role': 'user', 'content': user_content})
        if assistant_content:  # Only add the assistant's reply if it's not empty
            role_content_list.append({'role': 'assistant', 'content': assistant_content})
    return role_content_list

def role_content_to_history(role_content_list):
    history = {
        'internal': [],
        'visible': []
    }
    
    for i in range(0, len(role_content_list), 2):  # Iterate by twos to get user and (possible) assistant pairs
        user_content = role_content_list[i]['content']
        assistant_content = role_content_list[i+1]['content'] if (i+1 < len(role_content_list) and role_content_list[i+1]['role'] == 'assistant') else ''
        
        history['internal'].append([user_content, assistant_content])
        history['visible'].append([user_content, assistant_content])
    
    return history
ooba/utils/openai_messages_converters.py
KillianLucas/ooba