Coverage for tests / integration / test_pipeline.py: 100%

1import pytest 

2 

3from pattern_lens.activations import activations_main 

4from pattern_lens.figures import figures_main 

5 

6SAVE_PATH_BASE: str = "tests/.temp/pipeline" 

7PROMPTS_PATH: str = "data/pile_100.jsonl" 

8N_SAMPLES: int = 3 

9MIN_CHARS: int = 32 

10MAX_CHARS: int = 128 

11FORCE: bool = True 

12NO_INDEX_HTML: bool = False 

13FIGURES_PARALLEL: bool = False 

14 

15 

16@pytest.mark.parametrize( 

17 "model_name", 

18 ["pythia-14m", "tiny-stories-1M"], 

19) 

20def test_pipeline(model_name: str): 

21 activations_main( 

22 model_name=model_name, 

23 save_path=SAVE_PATH_BASE, 

24 prompts_path=PROMPTS_PATH, 

25 raw_prompts=True, 

26 min_chars=MIN_CHARS, 

27 max_chars=MAX_CHARS, 

28 force=FORCE, 

29 n_samples=N_SAMPLES, 

30 no_index_html=NO_INDEX_HTML, 

31 ) 

32 figures_main( 

33 model_name=model_name, 

34 save_path=SAVE_PATH_BASE, 

35 n_samples=N_SAMPLES, 

36 force=FORCE, 

37 parallel=FIGURES_PARALLEL, 

38 )