alpaca_lora_4bit/autograd_4bit/__init__.py

21 lines
720 B
Python

import os
from colorama import init, Fore, Back, Style
init(autoreset=True)
try:
GPTQ_VERSION = int(os.environ["GPTQ_VERSION"])
except:
print(Style.BRIGHT + Fore.YELLOW + "GPTQ_VERSION environment not provided. Fallback to GPTQv1")
GPTQ_VERSION = 1 # Fallback version
loader = None
if GPTQ_VERSION == 1:
from .autograd_4bit_v1 import Autograd4bitQuantLinear, load_llama_model_4bit_low_ram
print(Style.BRIGHT + Fore.GREEN + "GPTQv1 set")
elif GPTQ_VERSION == 2:
from .autograd_4bit_v2 import Autograd4bitQuantLinear, load_llama_model_4bit_low_ram
print(Style.BRIGHT + Fore.GREEN + "GPTQv2 set")
else:
raise ValueError("GPTQ_VERSION not set or invalid")