tests for additional merge fix unsloth zoo pr 163 (#2719)
* tests for additional merge fix unsloth zoo pr 163 * fixed load_dataset indent in mistral perplexity test file
This commit is contained in:
Roland Tannous
GitHub
parent
d535bf067e
commit
efe2cc43a7
18 changed files with 1094 additions and 11 deletions
|
|
@ -14,8 +14,10 @@ import gc
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
|
|
|||
|
|
@ -14,8 +14,9 @@ import gc
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
|
|
|||
|
|
@ -14,8 +14,10 @@ import gc
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
|
|
|||
|
|
@ -14,8 +14,9 @@ import gc
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
|
|
|||
|
|
@ -14,8 +14,9 @@ import gc
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
|
|
|||
|
|
@ -15,8 +15,9 @@ from huggingface_hub import HfFileSystem, hf_hub_download
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
|
|
|||
|
|
@ -15,8 +15,9 @@ from huggingface_hub import HfFileSystem, hf_hub_download
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.perplexity_eval import ppl_model, add_to_comparison, print_model_comparison
|
||||
|
|
|
|||
|
|
@ -12,8 +12,9 @@ from pathlib import Path
|
|||
import multiprocessing as mp
|
||||
import gc
|
||||
from multiprocessing import Queue
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.aime_eval import evaluate_model_aime, compare_aime_results
|
||||
|
|
|
|||
67
tests/saving/non_peft/test_mistral_non_peft.py
Normal file
67
tests/saving/non_peft/test_mistral_non_peft.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
from unsloth import FastLanguageModel
|
||||
from transformers import AutoModelForCausalLM
|
||||
from peft import PeftModel
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 PHASE 1: Loading Base Model")
|
||||
print(f"{'='*80}")
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name="unsloth/mistral-7b-v0.3",
|
||||
max_seq_length=2048,
|
||||
dtype=None,
|
||||
load_in_4bit=True,
|
||||
load_in_8bit=False,
|
||||
full_finetuning=False,
|
||||
)
|
||||
|
||||
|
||||
print("✅ Base model loaded successfully!")
|
||||
|
||||
### Attemtping save merge
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 PHASE 2: Attempting save_pretrained_merged (Should Warn)")
|
||||
print(f"{'='*80}")
|
||||
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
warnings.simplefilter("always")
|
||||
model.save_pretrained_merged("test_output", tokenizer)
|
||||
|
||||
# Verify warning
|
||||
assert len(w) >= 1, "Expected warning but none raised"
|
||||
warning_msg = str(w[0].message)
|
||||
expected_msg = "Model is not a PeftModel (no Lora adapters detected). Skipping Merge. Please use save_pretrained() or push_to_hub() instead!"
|
||||
assert expected_msg in warning_msg, f"Unexpected warning: {warning_msg}"
|
||||
assert expected_msg in warning_msg, f"Unexpected warning: {warning_msg}"
|
||||
|
||||
print("✅ Correct warning detected for non-PeftModel merge attempt!")
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 PHASE 3: Using save_pretrained (Should Succeed)")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
try:
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error") # Treat warnings as errors here
|
||||
model.save_pretrained("test_output")
|
||||
print("✅ Standard save_pretrained completed successfully!")
|
||||
except Exception as e:
|
||||
assert False, f"Phase 3 failed: {e}"
|
||||
|
||||
safe_remove_directory("./test_output")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
67
tests/saving/non_peft/test_whisper_non_peft.py
Normal file
67
tests/saving/non_peft/test_whisper_non_peft.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
from unsloth import FastLanguageModel, FastModel
|
||||
from transformers import AutoModelForCausalLM, WhisperForConditionalGeneration
|
||||
from peft import PeftModel
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 PHASE 1: Loading Base Model")
|
||||
print(f"{'='*80}")
|
||||
|
||||
model, tokenizer = FastModel.from_pretrained(
|
||||
model_name = "unsloth/whisper-large-v3",
|
||||
dtype = None, # Leave as None for auto detection
|
||||
load_in_4bit = False, # Set to True to do 4bit quantization which reduces memory
|
||||
auto_model = WhisperForConditionalGeneration,
|
||||
whisper_language = "English",
|
||||
whisper_task = "transcribe",
|
||||
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
|
||||
)
|
||||
|
||||
print("✅ Base model loaded successfully!")
|
||||
|
||||
### Attemtping save merge
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 PHASE 2: Attempting save_pretrained_merged (Should Warn)")
|
||||
print(f"{'='*80}")
|
||||
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
warnings.simplefilter("always")
|
||||
model.save_pretrained_merged("test_output", tokenizer)
|
||||
|
||||
# Verify warning
|
||||
assert len(w) >= 1, "Expected warning but none raised"
|
||||
warning_msg = str(w[0].message)
|
||||
expected_msg = "Model is not a PeftModel (no Lora adapters detected). Skipping Merge. Please use save_pretrained() or push_to_hub() instead!"
|
||||
assert expected_msg in warning_msg, f"Unexpected warning: {warning_msg}"
|
||||
assert expected_msg in warning_msg, f"Unexpected warning: {warning_msg}"
|
||||
|
||||
print("✅ Correct warning detected for non-PeftModel merge attempt!")
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 PHASE 3: Using save_pretrained (Should Succeed)")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
try:
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error") # Treat warnings as errors here
|
||||
model.save_pretrained("test_output")
|
||||
print("✅ Standard save_pretrained completed successfully!")
|
||||
except Exception as e:
|
||||
assert False, f"Phase 3 failed: {e}"
|
||||
|
||||
safe_remove_directory("./test_output")
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
156
tests/saving/text_to_speech_models/test_csm.py
Normal file
156
tests/saving/text_to_speech_models/test_csm.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
from unsloth import FastLanguageModel, FastModel
|
||||
from transformers import CsmForConditionalGeneration
|
||||
import torch
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from peft import PeftModel
|
||||
import warnings
|
||||
import requests
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.os_utils import require_package, require_python_package
|
||||
|
||||
require_package("ffmpeg", "ffmpeg")
|
||||
require_python_package("soundfile")
|
||||
|
||||
import soundfile as sf
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 1: Loading Model and LoRA Adapters")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
model, tokenizer = FastModel.from_pretrained(
|
||||
model_name = "unsloth/csm-1b",
|
||||
max_seq_length= 2048, # Choose any for long context!
|
||||
dtype = None, # Leave as None for auto-detection
|
||||
auto_model = CsmForConditionalGeneration,
|
||||
load_in_4bit = False, # Select True for 4bit - reduces memory usage
|
||||
)
|
||||
|
||||
|
||||
base_model_class = model.__class__.__name__
|
||||
|
||||
|
||||
model = FastModel.get_peft_model(
|
||||
model,
|
||||
r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",],
|
||||
lora_alpha = 32,
|
||||
lora_dropout = 0, # Supports any, but = 0 is optimized
|
||||
bias = "none", # Supports any, but = "none" is optimized
|
||||
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
|
||||
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
|
||||
random_state = 3407,
|
||||
use_rslora = False, # We support rank stabilized LoRA
|
||||
loftq_config = None, # And LoftQ
|
||||
)
|
||||
|
||||
print("✅ Model and LoRA adapters loaded successfully!")
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 2: Checking Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
|
||||
print("✅ Model is an instance of PeftModel!")
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 3: Checking Config Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
def find_lora_base_model(model_to_inspect):
|
||||
current = model_to_inspect
|
||||
if hasattr(current, "base_model"):
|
||||
current = current.base_model
|
||||
if hasattr(current, "model"):
|
||||
current = current.model
|
||||
return current
|
||||
pass
|
||||
|
||||
|
||||
config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
|
||||
|
||||
assert config_model.__class__.__name__ == base_model_class, f"Expected config_model class to be {base_model_class}"
|
||||
print("✅ config_model returns correct Base Model class:", str(base_model_class))
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 4: Saving and Merging Model")
|
||||
print(f"{'='*80}")
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error") # Treat warnings as errors
|
||||
try:
|
||||
model.save_pretrained_merged("csm", tokenizer)
|
||||
print("✅ Model saved and merged successfully without warnings!")
|
||||
except Exception as e:
|
||||
assert False, f"Model saving/merging failed with exception: {e}"
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 5: Loading Model for Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
model, processor = FastModel.from_pretrained(
|
||||
model_name = "./csm",
|
||||
max_seq_length= 2048, # Choose any for long context!
|
||||
dtype = None, # Leave as None for auto-detection
|
||||
auto_model = CsmForConditionalGeneration,
|
||||
load_in_4bit = False, # Select True for 4bit - reduces memory usage
|
||||
)
|
||||
|
||||
from transformers import AutoProcessor
|
||||
processor = AutoProcessor.from_pretrained("unsloth/csm-1b")
|
||||
|
||||
print("✅ Model loaded for inference successfully!")
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 6: Running Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
from transformers import pipeline
|
||||
import torch
|
||||
output_audio_path = "csm_audio.wav"
|
||||
try:
|
||||
text = "We just finished fine tuning a text to speech model... and it's pretty good!"
|
||||
speaker_id = 0
|
||||
inputs = processor(f"[{speaker_id}]{text}", add_special_tokens=True).to("cuda")
|
||||
audio_values = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=125, # 125 tokens is 10 seconds of audio, for longer speech increase this
|
||||
# play with these parameters to get the best results
|
||||
depth_decoder_temperature=0.6,
|
||||
depth_decoder_top_k=0,
|
||||
depth_decoder_top_p=0.9,
|
||||
temperature=0.8,
|
||||
top_k=50,
|
||||
top_p=1.0,
|
||||
#########################################################
|
||||
output_audio=True
|
||||
)
|
||||
audio = audio_values[0].to(torch.float32).cpu().numpy()
|
||||
sf.write("example_without_context.wav", audio, 24000)
|
||||
print(f"✅ Audio generated and saved to {output_audio_path}!")
|
||||
except Exception as e:
|
||||
assert False, f"Inference failed with exception: {e}"
|
||||
|
||||
|
||||
## assert that transcribed_text contains The birch canoe slid on the smooth planks. Glued the sheet to the dark blue background. It's easy to tell the depth of a well. Four hours of steady work faced us.
|
||||
|
||||
print("✅ All sections passed successfully!")
|
||||
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./csm")
|
||||
217
tests/saving/text_to_speech_models/test_lasa.py
Normal file
217
tests/saving/text_to_speech_models/test_lasa.py
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
from unsloth import FastLanguageModel, FastModel
|
||||
from transformers import CsmForConditionalGeneration
|
||||
import torch
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from peft import PeftModel
|
||||
import warnings
|
||||
import requests
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.os_utils import require_package, require_python_package
|
||||
|
||||
require_package("ffmpeg", "ffmpeg")
|
||||
require_python_package("soundfile")
|
||||
require_python_package("xcodec2")
|
||||
|
||||
import soundfile as sf
|
||||
from xcodec2.modeling_xcodec2 import XCodec2Model
|
||||
XCODEC2_MODEL_NAME = "HKUST-Audio/xcodec2"
|
||||
SAMPLE_RATE = 16000
|
||||
DEVICE = "cuda"
|
||||
|
||||
try:
|
||||
codec_model = XCodec2Model.from_pretrained(XCODEC2_MODEL_NAME)
|
||||
|
||||
except Exception as e:
|
||||
raise f"ERROR loading XCodec2 model: {e}."
|
||||
|
||||
codec_model.to('cpu')
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 1: Loading Model and LoRA Adapters")
|
||||
print(f"{'='*80}")
|
||||
|
||||
max_seq_length = 2048
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "unsloth/Llasa-1B",
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = None, # Select None for auto detection
|
||||
load_in_4bit = False, # Choose True for 4bit which reduces memory
|
||||
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
|
||||
)
|
||||
|
||||
base_model_class = model.__class__.__name__
|
||||
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r = 128, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
target_modules = ["q_proj", "v_proj"],
|
||||
lora_alpha = 128,
|
||||
lora_dropout = 0, # Supports any, but = 0 is optimized
|
||||
bias = "none", # Supports any, but = "none" is optimized
|
||||
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
|
||||
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
|
||||
random_state = 3407,
|
||||
use_rslora = False, # We support rank stabilized LoRA
|
||||
loftq_config = None, # And LoftQ
|
||||
)
|
||||
|
||||
print("✅ Model and LoRA adapters loaded successfully!")
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 2: Checking Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
|
||||
print("✅ Model is an instance of PeftModel!")
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 3: Checking Config Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
def find_lora_base_model(model_to_inspect):
|
||||
current = model_to_inspect
|
||||
if hasattr(current, "base_model"):
|
||||
current = current.base_model
|
||||
if hasattr(current, "model"):
|
||||
current = current.model
|
||||
return current
|
||||
pass
|
||||
|
||||
|
||||
config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
|
||||
|
||||
assert config_model.__class__.__name__ == base_model_class, f"Expected config_model class to be {base_model_class}"
|
||||
print("✅ config_model returns correct Base Model class:", str(base_model_class))
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 4: Saving and Merging Model")
|
||||
print(f"{'='*80}")
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error") # Treat warnings as errors
|
||||
try:
|
||||
model.save_pretrained_merged("lasa", tokenizer)
|
||||
print("✅ Model saved and merged successfully without warnings!")
|
||||
except Exception as e:
|
||||
assert False, f"Model saving/merging failed with exception: {e}"
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 5: Loading Model for Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "./lasa",
|
||||
max_seq_length = max_seq_length,
|
||||
dtype = None, # Select None for auto detection
|
||||
load_in_4bit = False, # Choose True for 4bit which reduces memory
|
||||
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
|
||||
)
|
||||
|
||||
#from transformers import AutoProcessor
|
||||
#processor = AutoProcessor.from_pretrained("unsloth/csm-1b")
|
||||
|
||||
print("✅ Model loaded for inference successfully!")
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 6: Running Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
from transformers import pipeline
|
||||
import torch
|
||||
output_audio_path = "lasa_audio.wav"
|
||||
input_text = "Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person."
|
||||
|
||||
FastLanguageModel.for_inference(model)
|
||||
|
||||
def ids_to_speech_tokens(speech_ids):
|
||||
|
||||
speech_tokens_str = []
|
||||
for speech_id in speech_ids:
|
||||
speech_tokens_str.append(f"<|s_{speech_id}|>")
|
||||
return speech_tokens_str
|
||||
|
||||
def extract_speech_ids(speech_tokens_str):
|
||||
|
||||
speech_ids = []
|
||||
for token_str in speech_tokens_str:
|
||||
if token_str.startswith('<|s_') and token_str.endswith('|>'):
|
||||
num_str = token_str[4:-2]
|
||||
|
||||
num = int(num_str)
|
||||
speech_ids.append(num)
|
||||
else:
|
||||
print(f"Unexpected token: {token_str}")
|
||||
return speech_ids
|
||||
|
||||
#TTS start!
|
||||
with torch.inference_mode():
|
||||
with torch.amp.autocast('cuda',dtype=model.dtype):
|
||||
formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
|
||||
|
||||
# Tokenize the text
|
||||
chat = [
|
||||
{"role": "user", "content": "Convert the text to speech:" + formatted_text},
|
||||
{"role": "assistant", "content": "<|SPEECH_GENERATION_START|>"}
|
||||
]
|
||||
|
||||
input_ids = tokenizer.apply_chat_template(
|
||||
chat,
|
||||
tokenize=True,
|
||||
return_tensors='pt',
|
||||
continue_final_message=True
|
||||
)
|
||||
input_ids = input_ids.to('cuda')
|
||||
|
||||
speech_end_id = tokenizer.convert_tokens_to_ids('<|SPEECH_GENERATION_END|>')
|
||||
|
||||
# Generate the speech autoregressively
|
||||
outputs = model.generate(
|
||||
input_ids,
|
||||
max_length=2048, # We trained our model with a max length of 2048
|
||||
eos_token_id= speech_end_id ,
|
||||
do_sample=True,
|
||||
top_p=1.2, # Adjusts the diversity of generated content
|
||||
temperature=1.2, # Controls randomness in output
|
||||
)
|
||||
# Extract the speech tokens
|
||||
generated_ids = outputs[0][input_ids.shape[1]:-1]
|
||||
|
||||
speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
|
||||
|
||||
# Convert token <|s_23456|> to int 23456
|
||||
speech_tokens = extract_speech_ids(speech_tokens)
|
||||
|
||||
speech_tokens = torch.tensor(speech_tokens).cpu().unsqueeze(0).unsqueeze(0)
|
||||
|
||||
# Decode the speech tokens to speech waveform
|
||||
gen_wav = codec_model.decode_code(speech_tokens)
|
||||
try:
|
||||
sf.write(output_audio_path, gen_wav[0, 0, :].cpu().numpy(), 16000)
|
||||
except Exception as e:
|
||||
assert False, f"Inference failed with exception: {e}"
|
||||
|
||||
|
||||
## assert that transcribed_text contains The birch canoe slid on the smooth planks. Glued the sheet to the dark blue background. It's easy to tell the depth of a well. Four hours of steady work faced us.
|
||||
|
||||
print("✅ All sections passed successfully!")
|
||||
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./lasa")
|
||||
254
tests/saving/text_to_speech_models/test_orpheus.py
Normal file
254
tests/saving/text_to_speech_models/test_orpheus.py
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
from unsloth import FastLanguageModel, FastModel
|
||||
from transformers import CsmForConditionalGeneration
|
||||
import torch
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from peft import PeftModel
|
||||
import warnings
|
||||
import requests
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.os_utils import require_package, require_python_package
|
||||
|
||||
require_package("ffmpeg", "ffmpeg")
|
||||
require_python_package("soundfile")
|
||||
require_python_package("snac")
|
||||
|
||||
import soundfile as sf
|
||||
from snac import SNAC
|
||||
snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
|
||||
snac_model = snac_model.to("cuda")
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 1: Loading Model and LoRA Adapters")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "unsloth/orpheus-3b-0.1-ft",
|
||||
max_seq_length= 2048, # Choose any for long context!
|
||||
dtype = None, # Select None for auto detection
|
||||
load_in_4bit = False, # Select True for 4bit which reduces memory usage
|
||||
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
|
||||
)
|
||||
|
||||
base_model_class = model.__class__.__name__
|
||||
|
||||
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r = 64, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
|
||||
"gate_proj", "up_proj", "down_proj",],
|
||||
lora_alpha = 64,
|
||||
lora_dropout = 0, # Supports any, but = 0 is optimized
|
||||
bias = "none", # Supports any, but = "none" is optimized
|
||||
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
|
||||
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
|
||||
random_state = 3407,
|
||||
use_rslora = False, # We support rank stabilized LoRA
|
||||
loftq_config = None, # And LoftQ
|
||||
)
|
||||
print("✅ Model and LoRA adapters loaded successfully!")
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 2: Checking Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
|
||||
print("✅ Model is an instance of PeftModel!")
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 3: Checking Config Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
def find_lora_base_model(model_to_inspect):
|
||||
current = model_to_inspect
|
||||
if hasattr(current, "base_model"):
|
||||
current = current.base_model
|
||||
if hasattr(current, "model"):
|
||||
current = current.model
|
||||
return current
|
||||
pass
|
||||
|
||||
|
||||
config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
|
||||
|
||||
assert config_model.__class__.__name__ == base_model_class, f"Expected config_model class to be {base_model_class}"
|
||||
print("✅ config_model returns correct Base Model class:", str(base_model_class))
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 4: Saving and Merging Model")
|
||||
print(f"{'='*80}")
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error") # Treat warnings as errors
|
||||
try:
|
||||
model.save_pretrained_merged("orpheus", tokenizer)
|
||||
print("✅ Model saved and merged successfully without warnings!")
|
||||
except Exception as e:
|
||||
assert False, f"Model saving/merging failed with exception: {e}"
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 5: Loading Model for Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name = "unsloth/orpheus-3b-0.1-ft",
|
||||
max_seq_length= 2048, # Choose any for long context!
|
||||
dtype = None, # Select None for auto detection
|
||||
load_in_4bit = False, # Select True for 4bit which reduces memory usage
|
||||
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
|
||||
)
|
||||
|
||||
#from transformers import AutoProcessor
|
||||
#processor = AutoProcessor.from_pretrained("unsloth/csm-1b")
|
||||
|
||||
print("✅ Model loaded for inference successfully!")
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 6: Running Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
#@title Run Inference
|
||||
|
||||
|
||||
FastLanguageModel.for_inference(model) # Enable native 2x faster inference
|
||||
|
||||
# Moving snac_model cuda to cpu
|
||||
snac_model.to("cpu")
|
||||
prompts = [
|
||||
"Hey there my name is Elise, <giggles> and I'm a speech generation model that can sound like a person.",
|
||||
]
|
||||
|
||||
chosen_voice = None # None for single-speaker
|
||||
|
||||
prompts_ = [(f"{chosen_voice}: " + p) if chosen_voice else p for p in prompts]
|
||||
|
||||
all_input_ids = []
|
||||
|
||||
for prompt in prompts_:
|
||||
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
|
||||
all_input_ids.append(input_ids)
|
||||
|
||||
start_token = torch.tensor([[ 128259]], dtype=torch.int64) # Start of human
|
||||
end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64) # End of text, End of human
|
||||
|
||||
all_modified_input_ids = []
|
||||
for input_ids in all_input_ids:
|
||||
modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1) # SOH SOT Text EOT EOH
|
||||
all_modified_input_ids.append(modified_input_ids)
|
||||
|
||||
all_padded_tensors = []
|
||||
all_attention_masks = []
|
||||
max_length = max([modified_input_ids.shape[1] for modified_input_ids in all_modified_input_ids])
|
||||
for modified_input_ids in all_modified_input_ids:
|
||||
padding = max_length - modified_input_ids.shape[1]
|
||||
padded_tensor = torch.cat([torch.full((1, padding), 128263, dtype=torch.int64), modified_input_ids], dim=1)
|
||||
attention_mask = torch.cat([torch.zeros((1, padding), dtype=torch.int64), torch.ones((1, modified_input_ids.shape[1]), dtype=torch.int64)], dim=1)
|
||||
all_padded_tensors.append(padded_tensor)
|
||||
all_attention_masks.append(attention_mask)
|
||||
|
||||
all_padded_tensors = torch.cat(all_padded_tensors, dim=0)
|
||||
all_attention_masks = torch.cat(all_attention_masks, dim=0)
|
||||
|
||||
input_ids = all_padded_tensors.to("cuda")
|
||||
attention_mask = all_attention_masks.to("cuda")
|
||||
generated_ids = model.generate(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
max_new_tokens=1200,
|
||||
do_sample=True,
|
||||
temperature=0.6,
|
||||
top_p=0.95,
|
||||
repetition_penalty=1.1,
|
||||
num_return_sequences=1,
|
||||
eos_token_id=128258,
|
||||
use_cache = True
|
||||
)
|
||||
token_to_find = 128257
|
||||
token_to_remove = 128258
|
||||
|
||||
token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)
|
||||
|
||||
if len(token_indices[1]) > 0:
|
||||
last_occurrence_idx = token_indices[1][-1].item()
|
||||
cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
|
||||
else:
|
||||
cropped_tensor = generated_ids
|
||||
|
||||
mask = cropped_tensor != token_to_remove
|
||||
|
||||
processed_rows = []
|
||||
|
||||
for row in cropped_tensor:
|
||||
masked_row = row[row != token_to_remove]
|
||||
processed_rows.append(masked_row)
|
||||
|
||||
code_lists = []
|
||||
|
||||
for row in processed_rows:
|
||||
row_length = row.size(0)
|
||||
new_length = (row_length // 7) * 7
|
||||
trimmed_row = row[:new_length]
|
||||
trimmed_row = [t - 128266 for t in trimmed_row]
|
||||
code_lists.append(trimmed_row)
|
||||
|
||||
|
||||
def redistribute_codes(code_list):
|
||||
layer_1 = []
|
||||
layer_2 = []
|
||||
layer_3 = []
|
||||
for i in range((len(code_list)+1)//7):
|
||||
layer_1.append(code_list[7*i])
|
||||
layer_2.append(code_list[7*i+1]-4096)
|
||||
layer_3.append(code_list[7*i+2]-(2*4096))
|
||||
layer_3.append(code_list[7*i+3]-(3*4096))
|
||||
layer_2.append(code_list[7*i+4]-(4*4096))
|
||||
layer_3.append(code_list[7*i+5]-(5*4096))
|
||||
layer_3.append(code_list[7*i+6]-(6*4096))
|
||||
codes = [torch.tensor(layer_1).unsqueeze(0),
|
||||
torch.tensor(layer_2).unsqueeze(0),
|
||||
torch.tensor(layer_3).unsqueeze(0)]
|
||||
|
||||
# codes = [c.to("cuda") for c in codes]
|
||||
audio_hat = snac_model.decode(codes)
|
||||
return audio_hat
|
||||
|
||||
my_samples = []
|
||||
for code_list in code_lists:
|
||||
samples = redistribute_codes(code_list)
|
||||
my_samples.append(samples)
|
||||
output_path = "orpheus_audio.wav"
|
||||
try:
|
||||
for i, samples in enumerate(my_samples):
|
||||
audio_data = samples.detach().squeeze().cpu().numpy()
|
||||
import soundfile as sf
|
||||
sf.write(output_path, audio_data, 24000) # Explicitly pass sample rate
|
||||
print(f"✅ Audio saved to {output_path}!")
|
||||
except Exception as e:
|
||||
assert False, f"Inference failed with exception: {e}"
|
||||
|
||||
# Verify the file exists
|
||||
import os
|
||||
assert os.path.exists(output_path), f"Audio file not found at {output_path}"
|
||||
print("✅ Audio file exists on disk!")
|
||||
del my_samples, samples
|
||||
## assert that transcribed_text contains The birch canoe slid on the smooth planks. Glued the sheet to the dark blue background. It's easy to tell the depth of a well. Four hours of steady work faced us.
|
||||
|
||||
print("✅ All sections passed successfully!")
|
||||
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./orpheus")
|
||||
189
tests/saving/text_to_speech_models/test_whisper.py
Normal file
189
tests/saving/text_to_speech_models/test_whisper.py
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
from unsloth import FastLanguageModel, FastModel
|
||||
from transformers import WhisperForConditionalGeneration, WhisperProcessor
|
||||
import torch
|
||||
# ruff: noqa
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from peft import PeftModel
|
||||
import warnings
|
||||
import requests
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.os_utils import require_package, require_python_package
|
||||
|
||||
require_package("ffmpeg", "ffmpeg")
|
||||
require_python_package("soundfile")
|
||||
|
||||
import soundfile as sf
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 1: Loading Model and LoRA Adapters")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
model, tokenizer = FastModel.from_pretrained(
|
||||
model_name = "unsloth/whisper-large-v3",
|
||||
dtype = None, # Leave as None for auto detection
|
||||
load_in_4bit = False, # Set to True to do 4bit quantization which reduces memory
|
||||
auto_model = WhisperForConditionalGeneration,
|
||||
whisper_language = "English",
|
||||
whisper_task = "transcribe",
|
||||
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
|
||||
)
|
||||
|
||||
|
||||
base_model_class = model.__class__.__name__
|
||||
#https://github.com/huggingface/transformers/issues/37172
|
||||
model.generation_config.input_ids = model.generation_config.forced_decoder_ids
|
||||
model.generation_config.forced_decoder_ids = None
|
||||
|
||||
|
||||
model = FastModel.get_peft_model(
|
||||
model,
|
||||
r = 64, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
|
||||
target_modules = ["q_proj", "v_proj"],
|
||||
lora_alpha = 64,
|
||||
lora_dropout = 0, # Supports any, but = 0 is optimized
|
||||
bias = "none", # Supports any, but = "none" is optimized
|
||||
# [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
|
||||
use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
|
||||
random_state = 3407,
|
||||
use_rslora = False, # We support rank stabilized LoRA
|
||||
loftq_config = None, # And LoftQ
|
||||
task_type = None, # ** MUST set this for Whisper **
|
||||
)
|
||||
|
||||
print("✅ Model and LoRA adapters loaded successfully!")
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 2: Checking Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
assert isinstance(model, PeftModel), "Model should be an instance of PeftModel"
|
||||
print("✅ Model is an instance of PeftModel!")
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 3: Checking Config Model Class Type")
|
||||
print(f"{'='*80}")
|
||||
|
||||
def find_lora_base_model(model_to_inspect):
|
||||
current = model_to_inspect
|
||||
if hasattr(current, "base_model"):
|
||||
current = current.base_model
|
||||
if hasattr(current, "model"):
|
||||
current = current.model
|
||||
return current
|
||||
pass
|
||||
|
||||
|
||||
config_model = find_lora_base_model(model) if isinstance(model, PeftModel) else model
|
||||
|
||||
assert config_model.__class__.__name__ == base_model_class, f"Expected config_model class to be {base_model_class}"
|
||||
print("✅ config_model returns correct Base Model class:", str(base_model_class))
|
||||
|
||||
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 4: Saving and Merging Model")
|
||||
print(f"{'='*80}")
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error") # Treat warnings as errors
|
||||
try:
|
||||
model.save_pretrained_merged("whisper", tokenizer)
|
||||
print("✅ Model saved and merged successfully without warnings!")
|
||||
except Exception as e:
|
||||
assert False, f"Model saving/merging failed with exception: {e}"
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 5: Loading Model for Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
model, tokenizer = FastModel.from_pretrained(
|
||||
model_name = "./whisper",
|
||||
dtype = None, # Leave as None for auto detection
|
||||
load_in_4bit = False, # Set to True to do 4bit quantization which reduces memory
|
||||
auto_model = WhisperForConditionalGeneration,
|
||||
whisper_language = "English",
|
||||
whisper_task = "transcribe",
|
||||
# token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
|
||||
)
|
||||
|
||||
# model = WhisperForConditionalGeneration.from_pretrained("./whisper")
|
||||
# processor = WhisperProcessor.from_pretrained("./whisper")
|
||||
|
||||
print("✅ Model loaded for inference successfully!")
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 6: Downloading Sample Audio File")
|
||||
print(f"{'='*80}")
|
||||
|
||||
audio_url = "https://upload.wikimedia.org/wikipedia/commons/5/5b/Speech_12dB_s16.flac"
|
||||
audio_file = "Speech_12dB_s16.flac"
|
||||
|
||||
try:
|
||||
headers = {
|
||||
"User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
|
||||
}
|
||||
response = requests.get(audio_url, headers=headers)
|
||||
response.raise_for_status()
|
||||
with open(audio_file, "wb") as f:
|
||||
f.write(response.content)
|
||||
print("✅ Audio file downloaded successfully!")
|
||||
except Exception as e:
|
||||
assert False, f"Failed to download audio file: {e}"
|
||||
|
||||
print(f"\n{'='*80}")
|
||||
print("🔍 SECTION 7: Running Inference")
|
||||
print(f"{'='*80}")
|
||||
|
||||
|
||||
from transformers import pipeline
|
||||
import torch
|
||||
FastModel.for_inference(model)
|
||||
model.eval()
|
||||
#Create pipeline without specifying the device
|
||||
whisper = pipeline(
|
||||
"automatic-speech-recognition",
|
||||
model=model,
|
||||
tokenizer=tokenizer.tokenizer,
|
||||
feature_extractor=tokenizer.feature_extractor,
|
||||
processor=tokenizer,
|
||||
return_language=True,
|
||||
torch_dtype=torch.float16 # Remove the device parameter
|
||||
)
|
||||
# Example usage
|
||||
audio_file = "Speech_12dB_s16.flac"
|
||||
transcribed_text = whisper(audio_file)
|
||||
# audio, sr = sf.read(audio_file)
|
||||
# input_features = processor(audio, return_tensors="pt").input_features
|
||||
# transcribed_text = model.generate(input_features=input_features)
|
||||
print(f"📝 Transcribed Text: {transcribed_text['text']}")
|
||||
|
||||
## assert that transcribed_text contains The birch canoe slid on the smooth planks. Glued the sheet to the dark blue background. It's easy to tell the depth of a well. Four hours of steady work faced us.
|
||||
|
||||
expected_phrases = [
|
||||
"birch canoe slid on the smooth planks",
|
||||
"sheet to the dark blue background",
|
||||
"easy to tell the depth of a well",
|
||||
"Four hours of steady work faced us",
|
||||
]
|
||||
|
||||
transcribed_lower = transcribed_text["text"].lower()
|
||||
all_phrases_found = all(phrase.lower() in transcribed_lower for phrase in expected_phrases)
|
||||
|
||||
assert all_phrases_found, f"Expected phrases not found in transcription: {transcribed_text['text']}"
|
||||
print("✅ Transcription contains all expected phrases!")
|
||||
|
||||
|
||||
safe_remove_directory("./unsloth_compiled_cache")
|
||||
safe_remove_directory("./whisper")
|
||||
|
|
@ -11,8 +11,9 @@ from huggingface_hub import HfFileSystem
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
|
||||
|
|
|
|||
|
|
@ -11,8 +11,10 @@ from trl import SFTTrainer, SFTConfig
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
|
||||
|
|
|
|||
|
|
@ -11,8 +11,9 @@ from trl import SFTTrainer, SFTConfig
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
REPO_ROOT = Path(__file__).parents[3]
|
||||
sys.path.append(str(REPO_ROOT))
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from tests.utils.cleanup_utils import safe_remove_directory
|
||||
from tests.utils.ocr_eval import OCRModelEvaluator
|
||||
|
|
|
|||
119
tests/utils/os_utils.py
Normal file
119
tests/utils/os_utils.py
Normal file
|
|
@ -0,0 +1,119 @@
|
|||
import subprocess
|
||||
import sys
|
||||
import os
|
||||
import shutil
|
||||
import importlib
|
||||
|
||||
def detect_package_manager():
|
||||
"""Detect the available package manager"""
|
||||
package_managers = {
|
||||
'apt': '/usr/bin/apt',
|
||||
'yum': '/usr/bin/yum',
|
||||
'dnf': '/usr/bin/dnf',
|
||||
'pacman': '/usr/bin/pacman',
|
||||
'zypper': '/usr/bin/zypper'
|
||||
}
|
||||
|
||||
for pm, path in package_managers.items():
|
||||
if os.path.exists(path):
|
||||
return pm
|
||||
return None
|
||||
|
||||
def check_package_installed(package_name, package_manager=None):
|
||||
"""Check if a package is installed using the system package manager"""
|
||||
|
||||
if package_manager is None:
|
||||
package_manager = detect_package_manager()
|
||||
|
||||
if package_manager is None:
|
||||
print("Warning: Could not detect package manager")
|
||||
return None
|
||||
|
||||
try:
|
||||
if package_manager == 'apt':
|
||||
# Check with dpkg
|
||||
result = subprocess.run(['dpkg', '-l', package_name],
|
||||
capture_output=True, text=True)
|
||||
return result.returncode == 0
|
||||
|
||||
elif package_manager in ['yum', 'dnf']:
|
||||
# Check with rpm
|
||||
result = subprocess.run(['rpm', '-q', package_name],
|
||||
capture_output=True, text=True)
|
||||
return result.returncode == 0
|
||||
|
||||
elif package_manager == 'pacman':
|
||||
result = subprocess.run(['pacman', '-Q', package_name],
|
||||
capture_output=True, text=True)
|
||||
return result.returncode == 0
|
||||
|
||||
elif package_manager == 'zypper':
|
||||
result = subprocess.run(['zypper', 'se', '-i', package_name],
|
||||
capture_output=True, text=True)
|
||||
return package_name in result.stdout
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error checking package: {e}")
|
||||
return None
|
||||
|
||||
def require_package(package_name, executable_name=None):
|
||||
"""Require a package to be installed, exit if not found"""
|
||||
|
||||
# First check if executable is in PATH (most reliable)
|
||||
if executable_name:
|
||||
if shutil.which(executable_name):
|
||||
print(f"✓ {executable_name} is available")
|
||||
return
|
||||
|
||||
# Then check with package manager
|
||||
pm = detect_package_manager()
|
||||
is_installed = check_package_installed(package_name, pm)
|
||||
|
||||
if is_installed:
|
||||
print(f"✓ Package {package_name} is installed")
|
||||
return
|
||||
|
||||
# Package not found - show installation instructions
|
||||
print(f"❌ Error: {package_name} is not installed")
|
||||
print(f"\nPlease install {package_name} using your system package manager:")
|
||||
|
||||
install_commands = {
|
||||
'apt': f"sudo apt update && sudo apt install {package_name}",
|
||||
'yum': f"sudo yum install {package_name}",
|
||||
'dnf': f"sudo dnf install {package_name}",
|
||||
'pacman': f"sudo pacman -S {package_name}",
|
||||
'zypper': f"sudo zypper install {package_name}"
|
||||
}
|
||||
|
||||
if pm and pm in install_commands:
|
||||
print(f" {install_commands[pm]}")
|
||||
else:
|
||||
for pm_name, cmd in install_commands.items():
|
||||
print(f" {pm_name}: {cmd}")
|
||||
|
||||
print(f"\nAlternatively, install with conda:")
|
||||
print(f" conda install -c conda-forge {package_name}")
|
||||
|
||||
print(f"\nPlease install the required package and run the script again.")
|
||||
sys.exit(1)
|
||||
|
||||
# Usage
|
||||
#require_package("ffmpeg", "ffmpeg")
|
||||
|
||||
def require_python_package(package_name, import_name=None, pip_name=None):
|
||||
"""Require a Python package to be installed, exit if not found"""
|
||||
if import_name is None:
|
||||
import_name = package_name
|
||||
if pip_name is None:
|
||||
pip_name = package_name
|
||||
|
||||
if importlib.util.find_spec(import_name) is None:
|
||||
print(f"❌ Error: Python package '{package_name}' is not installed")
|
||||
print(f"\nPlease install {package_name} using pip:")
|
||||
print(f" pip install {pip_name}")
|
||||
print(f" # or with conda:")
|
||||
print(f" conda install {pip_name}")
|
||||
print(f"\nAfter installation, run this script again.")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print(f"✓ Python package '{package_name}' is installed")
|
||||
Loading…
Add table
Add a link
Reference in a new issue