Spaces:
Sleeping
Sleeping
from src.model import load_model, configure_peft_model | |
import torch | |
def test_gpu_feature(): | |
# Your test code that needs a GPU | |
assert torch.cuda.is_available() | |
def test_load_model(): | |
model_name = "unsloth/Meta-Llama-3.1-8B" | |
model, tokenizer = load_model(model_name, 16, None, True, {'': 0}) | |
# Check that model and tokenizer are not None | |
assert model is not None | |
assert tokenizer is not None | |
# Check that model is on the correct device (e.g., GPU or CPU) | |
assert next(model.parameters()).device == torch.device('cuda:0'), "Model should be loaded on CUDA device" | |
# Check that the tokenizer is an instance of the correct class | |
assert hasattr(tokenizer, "encode"), "Tokenizer should have the 'encode' method" | |
def test_configure_peft_model(): | |
model_name = "unsloth/Meta-Llama-3.1-8B" | |
model, _ = load_model(model_name, 16, None, True, {'': 0}) | |
# Configure the PEFT model | |
peft_model = configure_peft_model(model, target_modules=["q_proj", "down_proj"]) | |
# Check that PEFT model is not None | |
assert peft_model is not None, "PEFT model should not be None" | |
# Check that the PEFT model has a forward method | |
assert hasattr(peft_model, "forward"), "PEFT model should have a 'forward' method" | |
# Ensure that PEFT model can perform a forward pass (check if no error is raised) | |
try: | |
dummy_input = torch.randint(0, 1000, (1, 16)) # Dummy input tensor | |
peft_model(dummy_input) | |
except Exception as e: | |
pytest.fail(f"PEFT model forward pass failed: {e}") |