From fbeaaf1531c74bb89adb3449cac1afa8989005e6 Mon Sep 17 00:00:00 2001 From: artitw Date: Wed, 9 Oct 2024 04:17:06 +0000 Subject: [PATCH] Verbosity and error handling --- setup.py | 2 +- text2text/assistant.py | 23 +- text2text/rag_assistant.py | 9 +- text2text/utils/ollama_install.sh | 370 ------------------------------ 4 files changed, 27 insertions(+), 377 deletions(-) delete mode 100644 text2text/utils/ollama_install.sh diff --git a/setup.py b/setup.py index f970a53..22066b3 100755 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name="text2text", - version="1.7.3", + version="1.7.4", author="artitw", author_email="artitw@gmail.com", description="Text2Text: Crosslingual NLP/G toolkit", diff --git a/text2text/assistant.py b/text2text/assistant.py index 47b410f..5a5a7fd 100644 --- a/text2text/assistant.py +++ b/text2text/assistant.py @@ -4,6 +4,7 @@ import subprocess import warnings +from tqdm.auto import tqdm from llama_index.llms.ollama import Ollama from llama_index.core.llms import ChatMessage @@ -57,30 +58,41 @@ def __del__(self): warnings.warn(str(e)) def load_model(self): + pbar = tqdm(total=6, desc='Model setup') if not ollama_version(): self.__del__() + pbar.update(1) return_code = os.system("sudo apt install -q -y lshw") if return_code != 0: raise Exception("Cannot install lshw.") + pbar.update(1) result = os.system( "curl -fsSL https://ollama.com/install.sh | sh" ) if result != 0: raise Exception("Cannot install ollama") + pbar.update(1) self.ollama_serve_proc = subprocess.Popen(["ollama", "serve"]) time.sleep(1) + pbar.update(1) if not ollama_version(): raise Exception("Cannot serve ollama") + pbar.update(1) + else: + pbar.update(5) result = ollama.pull(self.model_name) if result["status"] == "success": ollama_run_proc = subprocess.Popen(["ollama", "run", self.model_name]) + pbar.update(1) else: raise Exception(f"Did not pull {self.model_name}. Try restarting.") + + pbar.close() def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, schema=None, **kwargs): try: @@ -96,8 +108,15 @@ def chat_completion(self, messages=[{"role": "user", "content": "hello"}], strea return self.chat_completion(messages=messages, stream=stream, schema=schema, **kwargs) if schema: - msgs = [ChatMessage(**m) for m in messages] - return self.structured_client.as_structured_llm(schema).chat(messages=msgs).raw + try: + msgs = [ChatMessage(**m) for m in messages] + return self.structured_client.as_structured_llm(schema).chat(messages=msgs).raw + except Exception as e: + warnings.warn(str(e)) + warnings.warn(f"Schema extraction failed for {messages}") + default_schema = schema() + warnings.warn(f"Returning schema with default values: {vars(default_schema)}") + return default_schema return self.client.chat(model=self.model_name, messages=messages, stream=stream) def embed(self, texts): diff --git a/text2text/rag_assistant.py b/text2text/rag_assistant.py index 805d141..0f17ade 100644 --- a/text2text/rag_assistant.py +++ b/text2text/rag_assistant.py @@ -4,6 +4,7 @@ import warnings import urllib.parse +from tqdm.auto import tqdm from bs4 import BeautifulSoup def get_cleaned_html(url): @@ -32,7 +33,7 @@ def __init__(self, **kwargs): texts = kwargs.get("texts", []) urls = kwargs.get("urls", []) input_lines = [] - for u in urls: + for u in tqdm(urls, desc='Scrape URLs'): if is_valid_url(u): try: texts.append(get_cleaned_html(u)) @@ -42,7 +43,7 @@ def __init__(self, **kwargs): warnings.warn(f"Skipping invalid URL: {u}") if schema: - for t in texts: + for t in tqdm(texts, desc='Schema extraction'): fields = ", ".join(schema.model_fields.keys()) prompt = f'Extract {fields} from the following text:\n\n{t}' res = t2t.Assistant.chat_completion(self, [{"role": "user", "content": prompt}], schema=schema) @@ -57,6 +58,6 @@ def chat_completion(self, messages=[{"role": "user", "content": "hello"}], strea k = kwargs.get("k", 3) query = messages[-1]["content"] docs = self.index.retrieve([query], k=k)[0] - grounding_information = "\n\n".join(docs) + "\n\n" - messages[-1] = {"role": "user", "content": grounding_information+query} + grounding_information = "Ground your response based on the following information:\n\n" + "\n- ".join(docs) + messages[-1] = {"role": "user", "content": query + "\n\n" + grounding_information} return t2t.Assistant.chat_completion(self, messages=messages, stream=stream, schema=schema, **kwargs) diff --git a/text2text/utils/ollama_install.sh b/text2text/utils/ollama_install.sh deleted file mode 100644 index 9a3c9d1..0000000 --- a/text2text/utils/ollama_install.sh +++ /dev/null @@ -1,370 +0,0 @@ -#!/bin/sh -# This script installs Ollama on Linux. -# It detects the current operating system architecture and installs the appropriate version of Ollama. - -set -eu - -status() { echo ">>> $*" >&2; } -error() { echo "ERROR $*"; exit 1; } -warning() { echo "WARNING: $*"; } - -TEMP_DIR=$(mktemp -d) -cleanup() { rm -rf $TEMP_DIR; } -trap cleanup EXIT - -available() { command -v $1 >/dev/null; } -require() { - local MISSING='' - for TOOL in $*; do - if ! available $TOOL; then - MISSING="$MISSING $TOOL" - fi - done - - echo $MISSING -} - -[ "$(uname -s)" = "Linux" ] || error 'This script is intended to run on Linux only.' - -ARCH=$(uname -m) -case "$ARCH" in - x86_64) ARCH="amd64" ;; - aarch64|arm64) ARCH="arm64" ;; - *) error "Unsupported architecture: $ARCH" ;; -esac - -IS_WSL2=false - -KERN=$(uname -r) -case "$KERN" in - *icrosoft*WSL2 | *icrosoft*wsl2) IS_WSL2=true;; - *icrosoft) error "Microsoft WSL1 is not currently supported. Please use WSL2 with 'wsl --set-version 2'" ;; - *) ;; -esac - -VER_PARAM="${OLLAMA_VERSION:+?version=$OLLAMA_VERSION}" - -SUDO= -if [ "$(id -u)" -ne 0 ]; then - # Running as root, no need for sudo - if ! available sudo; then - error "This script requires superuser permissions. Please re-run as root." - fi - - SUDO="sudo" -fi - -NEEDS=$(require curl awk grep sed tee xargs) -if [ -n "$NEEDS" ]; then - status "ERROR: The following tools are required but missing:" - for NEED in $NEEDS; do - echo " - $NEED" - done - exit 1 -fi - -for BINDIR in /usr/local/bin /usr/bin /bin; do - echo $PATH | grep -q $BINDIR && break || continue -done -OLLAMA_INSTALL_DIR=$(dirname ${BINDIR}) - -status "Installing ollama to $OLLAMA_INSTALL_DIR" -$SUDO install -o0 -g0 -m755 -d $BINDIR -$SUDO install -o0 -g0 -m755 -d "$OLLAMA_INSTALL_DIR" -if curl -I --silent --fail --location "https://ollama.com/download/ollama-linux-${ARCH}.tgz${VER_PARAM}" >/dev/null ; then - status "Downloading Linux ${ARCH} bundle" - curl --fail --show-error --location --progress-bar \ - "https://ollama.com/download/ollama-linux-${ARCH}.tgz${VER_PARAM}" | \ - $SUDO tar -xzf - -C "$OLLAMA_INSTALL_DIR" - BUNDLE=1 - if [ "$OLLAMA_INSTALL_DIR/bin/ollama" != "$BINDIR/ollama" ] ; then - status "Making ollama accessible in the PATH in $BINDIR" - $SUDO ln -sf "$OLLAMA_INSTALL_DIR/ollama" "$BINDIR/ollama" - fi -else - status "Downloading Linux ${ARCH} CLI" - curl --fail --show-error --location --progress-bar -o "$TEMP_DIR/ollama"\ - "https://ollama.com/download/ollama-linux-${ARCH}${VER_PARAM}" - $SUDO install -o0 -g0 -m755 $TEMP_DIR/ollama $OLLAMA_INSTALL_DIR/ollama - BUNDLE=0 - if [ "$OLLAMA_INSTALL_DIR/ollama" != "$BINDIR/ollama" ] ; then - status "Making ollama accessible in the PATH in $BINDIR" - $SUDO ln -sf "$OLLAMA_INSTALL_DIR/ollama" "$BINDIR/ollama" - fi -fi - - -install_success() { - status 'The Ollama API is now available at 127.0.0.1:11434.' - status 'Install complete. Run "ollama" from the command line.' -} -trap install_success EXIT - -# Everything from this point onwards is optional. - -configure_systemd() { - if ! id ollama >/dev/null 2>&1; then - status "Creating ollama user..." - $SUDO useradd -r -s /bin/false -U -m -d /usr/share/ollama ollama - fi - if getent group render >/dev/null 2>&1; then - status "Adding ollama user to render group..." - $SUDO usermod -a -G render ollama - fi - if getent group video >/dev/null 2>&1; then - status "Adding ollama user to video group..." - $SUDO usermod -a -G video ollama - fi - - status "Adding current user to ollama group..." - $SUDO usermod -a -G ollama $(whoami) - - status "Creating ollama systemd service..." - cat </dev/null -[Unit] -Description=Ollama Service -After=network-online.target - -[Service] -ExecStart=$BINDIR/ollama serve -User=ollama -Group=ollama -Restart=always -RestartSec=3 -Environment="PATH=$PATH" - -[Install] -WantedBy=default.target -EOF - SYSTEMCTL_RUNNING="$(systemctl is-system-running || true)" - case $SYSTEMCTL_RUNNING in - running|degraded) - status "Enabling and starting ollama service..." - $SUDO systemctl daemon-reload - $SUDO systemctl enable ollama - - start_service() { $SUDO systemctl restart ollama; } - trap start_service EXIT - ;; - esac -} - -if available systemctl; then - configure_systemd -fi - -# WSL2 only supports GPUs via nvidia passthrough -# so check for nvidia-smi to determine if GPU is available -if [ "$IS_WSL2" = true ]; then - if available nvidia-smi && [ -n "$(nvidia-smi | grep -o "CUDA Version: [0-9]*\.[0-9]*")" ]; then - status "Nvidia GPU detected." - fi - install_success - exit 0 -fi - -# Install GPU dependencies on Linux -if ! available lspci && ! available lshw; then - warning "Unable to detect NVIDIA/AMD GPU. Install lspci or lshw to automatically detect and install GPU dependencies." - exit 0 -fi - -check_gpu() { - # Look for devices based on vendor ID for NVIDIA and AMD - case $1 in - lspci) - case $2 in - nvidia) available lspci && lspci -d '10de:' | grep -q 'NVIDIA' || return 1 ;; - amdgpu) available lspci && lspci -d '1002:' | grep -q 'AMD' || return 1 ;; - esac ;; - lshw) - case $2 in - nvidia) available lshw && $SUDO lshw -c display -numeric -disable network | grep -q 'vendor: .* \[10DE\]' || return 1 ;; - amdgpu) available lshw && $SUDO lshw -c display -numeric -disable network | grep -q 'vendor: .* \[1002\]' || return 1 ;; - esac ;; - nvidia-smi) available nvidia-smi || return 1 ;; - esac -} - -if check_gpu nvidia-smi; then - status "NVIDIA GPU installed." - exit 0 -fi - -if ! check_gpu lspci nvidia && ! check_gpu lshw nvidia && ! check_gpu lspci amdgpu && ! check_gpu lshw amdgpu; then - install_success - warning "No NVIDIA/AMD GPU detected. Ollama will run in CPU-only mode." - exit 0 -fi - -if check_gpu lspci amdgpu || check_gpu lshw amdgpu; then - if [ $BUNDLE -ne 0 ]; then - status "Downloading Linux ROCm ${ARCH} bundle" - curl --fail --show-error --location --progress-bar \ - "https://ollama.com/download/ollama-linux-${ARCH}-rocm.tgz${VER_PARAM}" | \ - $SUDO tar -xzf - -C "$OLLAMA_INSTALL_DIR" - - install_success - status "AMD GPU ready." - exit 0 - fi - # Look for pre-existing ROCm v6 before downloading the dependencies - for search in "${HIP_PATH:-''}" "${ROCM_PATH:-''}" "/opt/rocm" "/usr/lib64"; do - if [ -n "${search}" ] && [ -e "${search}/libhipblas.so.2" -o -e "${search}/lib/libhipblas.so.2" ]; then - status "Compatible AMD GPU ROCm library detected at ${search}" - install_success - exit 0 - fi - done - - status "Downloading AMD GPU dependencies..." - $SUDO rm -rf /usr/share/ollama/lib - $SUDO chmod o+x /usr/share/ollama - $SUDO install -o ollama -g ollama -m 755 -d /usr/share/ollama/lib/rocm - curl --fail --show-error --location --progress-bar "https://ollama.com/download/ollama-linux-amd64-rocm.tgz${VER_PARAM}" \ - | $SUDO tar zx --owner ollama --group ollama -C /usr/share/ollama/lib/rocm . - install_success - status "AMD GPU ready." - exit 0 -fi - -CUDA_REPO_ERR_MSG="NVIDIA GPU detected, but your OS and Architecture are not supported by NVIDIA. Please install the CUDA driver manually https://docs.nvidia.com/cuda/cuda-installation-guide-linux/" -# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#rhel-7-centos-7 -# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#rhel-8-rocky-8 -# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#rhel-9-rocky-9 -# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#fedora -install_cuda_driver_yum() { - status 'Installing NVIDIA repository...' - - case $PACKAGE_MANAGER in - yum) - $SUDO $PACKAGE_MANAGER -y install yum-utils - if curl -I --silent --fail --location "https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo" >/dev/null ; then - $SUDO $PACKAGE_MANAGER-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo - else - error $CUDA_REPO_ERR_MSG - fi - ;; - dnf) - if curl -I --silent --fail --location "https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo" >/dev/null ; then - $SUDO $PACKAGE_MANAGER config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-$1$2.repo - else - error $CUDA_REPO_ERR_MSG - fi - ;; - esac - - case $1 in - rhel) - status 'Installing EPEL repository...' - # EPEL is required for third-party dependencies such as dkms and libvdpau - $SUDO $PACKAGE_MANAGER -y install https://dl.fedoraproject.org/pub/epel/epel-release-latest-$2.noarch.rpm || true - ;; - esac - - status 'Installing CUDA driver...' - - if [ "$1" = 'centos' ] || [ "$1$2" = 'rhel7' ]; then - $SUDO $PACKAGE_MANAGER -y install nvidia-driver-latest-dkms - fi - - $SUDO $PACKAGE_MANAGER -y install cuda-drivers -} - -# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#ubuntu -# ref: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#debian -install_cuda_driver_apt() { - status 'Installing NVIDIA repository...' - if curl -I --silent --fail --location "https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-keyring_1.1-1_all.deb" >/dev/null ; then - curl -fsSL -o $TEMP_DIR/cuda-keyring.deb https://developer.download.nvidia.com/compute/cuda/repos/$1$2/$(uname -m | sed -e 's/aarch64/sbsa/')/cuda-keyring_1.1-1_all.deb - else - error $CUDA_REPO_ERR_MSG - fi - - case $1 in - debian) - status 'Enabling contrib sources...' - $SUDO sed 's/main/contrib/' < /etc/apt/sources.list | $SUDO tee /etc/apt/sources.list.d/contrib.list > /dev/null - if [ -f "/etc/apt/sources.list.d/debian.sources" ]; then - $SUDO sed 's/main/contrib/' < /etc/apt/sources.list.d/debian.sources | $SUDO tee /etc/apt/sources.list.d/contrib.sources > /dev/null - fi - ;; - esac - - status 'Installing CUDA driver...' - $SUDO dpkg -i $TEMP_DIR/cuda-keyring.deb - $SUDO apt-get update - - [ -n "$SUDO" ] && SUDO_E="$SUDO -E" || SUDO_E= - DEBIAN_FRONTEND=noninteractive $SUDO_E apt-get -y install cuda-drivers -q -} - -if [ ! -f "/etc/os-release" ]; then - error "Unknown distribution. Skipping CUDA installation." -fi - -. /etc/os-release - -OS_NAME=$ID -OS_VERSION=$VERSION_ID - -PACKAGE_MANAGER= -for PACKAGE_MANAGER in dnf yum apt-get; do - if available $PACKAGE_MANAGER; then - break - fi -done - -if [ -z "$PACKAGE_MANAGER" ]; then - error "Unknown package manager. Skipping CUDA installation." -fi - -if ! check_gpu nvidia-smi || [ -z "$(nvidia-smi | grep -o "CUDA Version: [0-9]*\.[0-9]*")" ]; then - case $OS_NAME in - centos|rhel) install_cuda_driver_yum 'rhel' $(echo $OS_VERSION | cut -d '.' -f 1) ;; - rocky) install_cuda_driver_yum 'rhel' $(echo $OS_VERSION | cut -c1) ;; - fedora) [ $OS_VERSION -lt '39' ] && install_cuda_driver_yum $OS_NAME $OS_VERSION || install_cuda_driver_yum $OS_NAME '39';; - amzn) install_cuda_driver_yum 'fedora' '37' ;; - debian) install_cuda_driver_apt $OS_NAME $OS_VERSION ;; - ubuntu) install_cuda_driver_apt $OS_NAME $(echo $OS_VERSION | sed 's/\.//') ;; - *) exit ;; - esac -fi - -if ! lsmod | grep -q nvidia || ! lsmod | grep -q nvidia_uvm; then - KERNEL_RELEASE="$(uname -r)" - case $OS_NAME in - rocky) $SUDO $PACKAGE_MANAGER -y install kernel-devel kernel-headers ;; - centos|rhel|amzn) $SUDO $PACKAGE_MANAGER -y install kernel-devel-$KERNEL_RELEASE kernel-headers-$KERNEL_RELEASE ;; - fedora) $SUDO $PACKAGE_MANAGER -y install kernel-devel-$KERNEL_RELEASE ;; - debian|ubuntu) $SUDO apt-get -y install linux-headers-$KERNEL_RELEASE ;; - *) exit ;; - esac - - NVIDIA_CUDA_VERSION=$($SUDO dkms status | awk -F: '/added/ { print $1 }') - if [ -n "$NVIDIA_CUDA_VERSION" ]; then - $SUDO dkms install $NVIDIA_CUDA_VERSION - fi - - if lsmod | grep -q nouveau; then - status 'Reboot to complete NVIDIA CUDA driver install.' - exit 0 - fi - - $SUDO modprobe nvidia - $SUDO modprobe nvidia_uvm -fi - -# make sure the NVIDIA modules are loaded on boot with nvidia-persistenced -if available nvidia-persistenced; then - $SUDO touch /etc/modules-load.d/nvidia.conf - MODULES="nvidia nvidia-uvm" - for MODULE in $MODULES; do - if ! grep -qxF "$MODULE" /etc/modules-load.d/nvidia.conf; then - echo "$MODULE" | $SUDO tee -a /etc/modules-load.d/nvidia.conf > /dev/null - fi - done -fi - -status "NVIDIA GPU ready." -install_success \ No newline at end of file