test_lora.py 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. import pytest
  2. import os
  3. from utils import *
  4. server = ServerPreset.stories15m_moe()
  5. LORA_FILE_URL = "https://huggingface.co/ggml-org/stories15M_MOE/resolve/main/moe_shakespeare15M.gguf"
  6. @pytest.fixture(scope="module", autouse=True)
  7. def create_server():
  8. global server
  9. server = ServerPreset.stories15m_moe()
  10. # download lora file if needed
  11. file_name = LORA_FILE_URL.split('/').pop()
  12. lora_file = f'../../../{file_name}'
  13. if not os.path.exists(lora_file):
  14. print(f"Downloading {LORA_FILE_URL} to {lora_file}")
  15. with open(lora_file, 'wb') as f:
  16. f.write(requests.get(LORA_FILE_URL).content)
  17. print(f"Done downloading lora file")
  18. server.lora_files = [lora_file]
  19. @pytest.mark.parametrize("scale,re_content", [
  20. # without applying lora, the model should behave like a bedtime story generator
  21. (0.0, "(little|girl|three|years|old)+"),
  22. # with lora, the model should behave like a Shakespearean text generator
  23. (1.0, "(eye|love|glass|sun)+"),
  24. ])
  25. def test_lora(scale: float, re_content: str):
  26. global server
  27. server.start()
  28. res_lora_control = server.make_request("POST", "/lora-adapters", data=[
  29. {"id": 0, "scale": scale}
  30. ])
  31. assert res_lora_control.status_code == 200
  32. res = server.make_request("POST", "/completion", data={
  33. "prompt": "Look in thy glass",
  34. })
  35. assert res.status_code == 200
  36. assert match_regex(re_content, res.body["content"])