phi-2

mjc_notes.md

@@ -1515,7 +1515,7 @@ counterfactuals?
 - [x] what about nonlinear? nope
 
 Next
-- [ ] with noise on embeddings? this would allow large models agian. I'm really struggling with these small models!?
+- [ ] with noise on embeddings? this would allow large models agian. GI'm really struggling with these small models!?
 
 # 2023-09-22 13:36:27
 
@@ -1858,13 +1858,9 @@ QC
 
 - [x] I have poor choice coverag however
 maybe I should just use the huggingface chat template? https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1 
-
-
 - [ ] Also I may be intervening to much, as I get nonsensicle answers. I need to add the intervention QC!
-
-- [x] Why does wizard-vicuna not work... it's because of the stupid <s> token I think. It wasn't trained with it and it drops context
-
-- [x] Oh it looks like my few shots might be wrong?? - fixes
+- [x] Why does wizard-vicuna not work... it's because of the stupid `<`s`>` token I think. It wasn't trained with it and it drops context
+- [x] Oh it looks like my few shots might be wrong? - fixes
 
 # 2023-10-28 12:33:27
 
@@ -1940,11 +1936,12 @@ Right now I'm using pipelines that do an intervention.
 
 # Running
 
-```sh
-python notebooks/make_dataset2.py --max_examples 1720 220 --datasets imdb glue:qnli super_glue:boolq 
+~~~sh
+python notebooks/make_dataset2.py --max_examples 1720 220 --datasets imdb glue:qnli super_glue:boolq
+~~~
 
-```
-an in intervene.py/create_cache_interventions we get the activations that are used to intervent and get a pair of hidden states
+
+an in intervene.py/create_cache_interventions we get the activations that are used to intervene and get a pair of hidden states
 - rep_reading_pipeline.get_directions which uses PCA to get an intervention
 
 
@@ -2072,20 +2069,75 @@ Wow that's a lot. If I want to focus on just lying, I might need to focus on not
 # 2023-12-09 11:34:11
 
 Questions:
-- understand HALOs https://twitter.com/ethayarajh/status/1732837520784957476 https://github.com/ContextualAI/HALOs
+- [x] understand HALOs https://twitter.com/ethayarajh/status/1732837520784957476 https://github.com/ContextualAI/HALOs
   - so it's just DPO with a differen't activation function on the reward, and notably it can use reward text instead of ranked pairs, letting you skip SFT. In a way it's just SFT?
  - [ ] discrete states, just look up QVAE?
    - [ ] hmm some use a categorical, and the gumbel reparam trick for end to end backprop
    - [ ] some use VQ-VAE which I haven't looked at before but look promising. But I want to!
  - [ ] does my pcr probe work ok? how to debug? I guess I need to check acc from it for a start
- - [ ] does my conv vae work? maybe I need transposed conv blocks?
+ - [x] does my conv vae work? maybe I need transposed conv blocks? nah
      - perhaps just use https://github.com/ctallec/world-models/blob/master/models/vae.py#L10
      - perhaps I need to focus on important features? Or on a task?
        - e.g. if doing inference on the reconstructed parts, can I get the same output? (RAM heavy)
        - if just apply an importance multipier
+- [ ] see if the use og SGB here gives me usefull ideas https://www.lesswrong.com/posts/7fxusXdkMNmAhkAfc/finding-sparse-linear-connections-between-features-in-llms
 
 
 
-```py
-QVAE psuedocode
-```
+
+# 2023-12-13 21:51:56
+
+TODO:
+- [ ] mean diff prob
+  - [ ] **and check the probe!**
+- [ ] handle lots of HS
+  - For this we need to load as a stream
+- [ ] model that can lie?
+  - [ ] QLoRa
+  - [ ] Read Cognitive Dissonance: Why Do Language Model Outputs Disagree with Internal Representations of Truthfulness?
+  - [ ] better model? Phi-2?
+- [ ] variation on my probe.... maybe I don't need ranking if I have a VAE
+- [ ] maybe I can use my intervention as an importance matrix for the VAE loss?
+- [ ] maybe I can do that codebook VAE? How big?
+
+OK I can't load all my hs into mem, that's not ideal....
+
+and I tried the mean diff intervention since eluther like it, but the truthfull llama one seems broken
+
+# probes!
+
+How to sanity check them?
+- how well can a linear probe do compared to random?
+
+So I'm using the ones from https://github.dev/andyzoujm/representation-engineering/tree/main/repe_eval/examples/decoder_repe_eval.ipynb but maybe I should use 
+- [Eleuther](https://github.com/EleutherAI/concept-erasure) 
+  - huh this is weird. it does sgd on action. has a whitenessing matrix
+  - > Intuitively, LEACE de-means and whitens x, projects onto the subspace responsible for correlations between X and Z, then unwhitens the result. Finally, it subtracts this value from x, thereby surgically removing the linearly available information about Z.
+- or [honestllama](https://github.com/likenneth/honest_llama/blob/master/utils.py#L730)
+  - > We compare three different directions for the ITI activation shift. Probe Weight Direction is the direction found by linear probing in subsection 3.2. Intervening in this direction is equivalent to doing one gradient descent step on the head activation to maximize its probability of being predicted as truthful. Mass Mean Shift works by first calculating the average of truthful and false activations and then using the vector pointing from the false mean to the truthful mean for intervention. As a baseline, we also apply the Contrast-Consistent Search (CCS) technique, where the direction is found while only knowing pairwise information of internal activations (Burns et al., 2022).
+  - ![Honest LLAMA Table 3](img/2023-12-14-11-51-32.png)
+    - **Mean Mass shift was the best by far** (PCA not considered)
+  - For magnitude they get `direction = direction / np.linalg.norm(direction)` and `std(activations @ directions)` for each head
+  - [center of mass directions](https://github.dev/likenneth/honest_llama/blob/207bb14b2c005e0593487cca8d22e072cbcb987b/utils.py#L730)
+  - use_random_dir
+  - [linear regression.coef_](https://github.dev/likenneth/honest_llama/blob/207bb14b2c005e0593487cca8d22e072cbcb987b/utils.py#L644) which is m from `mx+c`
+- or the geometry-of-truth https://github.com/saprmarks/geometry-of-truth/blob/main/interventions.ipynb
+  - LRProbe - trained linear layer with sigmoid. Directon from the weight
+  - MMProbe - mass
+    - `direction = pos_mean - neg_mean`
+    - `covariance = centered_data.t() @ centered_data / acts.shape[0]`
+  - CCSProbe: this must be the clustering one from constrastive clustering, trained with Adam
+- or representation engineering
+  - > Among these models, both unsupervised methods such as PCA and K-Means, as well as the supervised technique of Mean Difference, consistently exhibit robust overall performance
+  - magnitudes?
+
+
+## datasets
+
+From honest_llama
+- tqa_mc2 (multi choice)
+- tca_gen (generation)
+- tqa_gen_end_q (generation)
+
+
+Looks like I need hidden states

notebooks/027_train_mse_bigger.ipynb

@@ -5,7 +5,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# distance and direciton\n",
+    "# distance and direction\n",
     "\n",
     "Let try to opt for distance and direction with\n",
     "\n",

notebooks/make_dataset2.py

@@ -1,7 +1,7 @@
 # %% [markdown]
 # Use pipelines as this https://github.com/wassname/representation-engineering/blob/random_comments_ignore/examples/honesty/honesty.ipynb
-# 
-# 
+#
+#
 
 # %%
 # import your package
@@ -13,9 +13,11 @@ from src.helpers.torch import clear_mem
 import numpy as np
 import pandas as pd
 from matplotlib import pyplot as plt
-plt.style.use('ggplot')
+
+plt.style.use("ggplot")
 
 import os, psutil
+
 max_dataset_memory = f"{psutil.virtual_memory().total //2}"
 os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = max_dataset_memory
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -23,6 +25,7 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 from pathlib import Path
 from tqdm.auto import tqdm
 from loguru import logger
+
 # logger.add(os.sys.stderr, format="{time} {level} {message}", level="INFO")
 logger.add("logs/make_dataset_{time}.log")
 
@@ -39,6 +42,7 @@ from simple_parsing import ArgumentParser
 import transformers
 from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM
 from src.repe import repe_pipeline_registry
+
 repe_pipeline_registry()
 
 from src.models.load import load_model
@@ -48,12 +52,13 @@ import pickle
 
 from src.config import root_folder
 import json
+
 # from datasets import Dataset, DatasetInfo
 import datasets
 from src.config import root_folder
 from pathvalidate import sanitize_filename
 from src.helpers.ds import ds_keep_cols
-from src.datasets.intervene import create_cache_interventions 
+from src.datasets.intervene import create_cache_interventions
 
 # from sklearn.linear_model import LogisticRegression
 # from sklearn.metrics import f1_score, roc_auc_score, accuracy_score
@@ -63,7 +68,6 @@ from src.datasets.intervene import create_cache_interventions
 # %%
 
 TEST = False
-batch_size = 2
 
 parser = ArgumentParser(add_help=False)
 parser.add_arguments(ExtractConfig, dest="run")
@@ -71,32 +75,56 @@ args = parser.parse_args()
 cfg = args.run
 print(cfg)
 
+batch_size = cfg.batch_size
+
 model, tokenizer = load_model(cfg.model)
 
 
-tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True)
+tokenizer_args = dict(
+    padding="max_length",
+    max_length=cfg.max_length,
+    truncation=True,
+    add_special_tokens=True,
+)
 
 # %%
 if TEST:
     # # cache busting for the transformers map and ds steps
     import shutil
-    shutil.rmtree('~/.cache/huggingface/datasets/generator')
+
+    shutil.rmtree("~/.cache/huggingface/datasets/generator")
 
 
 # %% [markdown]
 # # Intervention fit/load
 
-# %%
 
 # %%
 
-# N_fit_examples = 20
-N_fit_examples = 30
+# Fit an intervention
+# N_fit_examples = 60
 rep_token = -1
 
-honesty_rep_reader1 = create_cache_interventions(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size, rep_token=rep_token)
+honesty_rep_reader1 = create_cache_interventions(
+    model,
+    tokenizer,
+    cfg,
+    direction_method=cfg.intervention_direction_method,
+    N_fit_examples=cfg.intervention_fit_examples,
+    batch_size=batch_size,
+    rep_token=rep_token,
+)
 
-honesty_rep_reader2 = create_cache_interventions(model, tokenizer, cfg, N_fit_examples=N_fit_examples, batch_size=batch_size, rep_token=rep_token, get_negative=True)
+honesty_rep_reader2 = create_cache_interventions(
+    model,
+    tokenizer,
+    cfg,
+    direction_method=cfg.intervention_direction_method,
+    N_fit_examples=cfg.intervention_fit_examples,
+    batch_size=batch_size,
+    rep_token=rep_token,
+    get_negative=True,
+)
 
 hidden_layers = sorted(honesty_rep_reader1.directions.keys())
 hidden_layers
@@ -112,38 +140,44 @@ hidden_layers
 # # Control helpers
 
 
-
-
 # %% [markdown]
 # # Control
 
 # %%
 
 rep_control_pipeline2 = pipeline(
-    "rep-control2", 
-    model=model, 
-    tokenizer=tokenizer, 
-    layers=hidden_layers, 
-    max_length=cfg.max_length,)
+    "rep-control2",
+    model=model,
+    tokenizer=tokenizer,
+    layers=hidden_layers,
+    max_length=cfg.max_length,
+)
 rep_control_pipeline2
 
 
 # %%
 from src.datasets.intervene import get_activations_from_reader
 
-activations1 = get_activations_from_reader(honesty_rep_reader1, hidden_layers, dtype=model.dtype, device=model.device)
-activations2 = get_activations_from_reader(honesty_rep_reader2, hidden_layers, dtype=model.dtype, device=model.device)
+activations1 = get_activations_from_reader(
+    honesty_rep_reader1, hidden_layers, dtype=model.dtype, device=model.device
+)
+activations2 = get_activations_from_reader(
+    honesty_rep_reader2, hidden_layers, dtype=model.dtype, device=model.device
+)
 
 # %%
 
 
-
-
 # %%
 if TEST:
     # unit test pipeline: with multiple input types: single, list, generator, dataset
     ## single
-    input_types = {'single':dataset_train[0], 'list':[dataset_train[i] for i in range(3)], 'generator':iter(dataset_train.select(range(3))), 'dataset':dataset_train.select(range(3)).to_iterable_dataset()}
+    input_types = {
+        "single": dataset_train[0],
+        "list": [dataset_train[i] for i in range(3)],
+        "generator": iter(dataset_train.select(range(3))),
+        "dataset": dataset_train.select(range(3)).to_iterable_dataset(),
+    }
     for name, ds in input_types.items():
         print(f"==== {name} ====")
         r = rep_control_pipeline2(ds, activations=activations1, batch_size=2)
@@ -154,8 +188,7 @@ if TEST:
         else:
             r = list(r)
         print(f"Control: {len(r)}")
-        print(r[0]['input_ids'].shape)
-        
+        print(r[0]["input_ids"].shape)
 
 
 # %%
@@ -164,39 +197,73 @@ if TEST:
 from src.datasets.intervene import test_intervention_quality
 
 if TEST:
-    test_intervention_quality(dataset_train, activations1, model, rep_control_pipeline2, batch_size=batch_size)
-    
-    test_intervention_quality(dataset_train, activations2, model, rep_control_pipeline2, batch_size=batch_size)
+    test_intervention_quality(
+        dataset_train, activations1, model, rep_control_pipeline2, batch_size=batch_size
+    )
+
+    test_intervention_quality(
+        dataset_train, activations2, model, rep_control_pipeline2, batch_size=batch_size
+    )
 
 # %%
 
 
-def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch_size=2, split_type="train", debug=TEST):
-    "create a dataset of hidden states."""
-    
+def create_hs_ds(
+    ds_name,
+    ds_tokens,
+    pipeline,
+    activations=None,
+    f=None,
+    batch_size=2,
+    split_type="train",
+    debug=TEST,
+):
+    "create a dataset of hidden states." ""
+
     N = len(ds_tokens)
-    dataset_name = sanitize_filename(f"{cfg.model}_{ds_name}_{split_type}_{N}", replacement_text="_")
-    f = str(root_folder / '.ds'/ f"{dataset_name}")
-    logger.info(f"Creating dataset {dataset_name} with {len(ds_tokens)} examples at `{f}`")
-    
-    info_kwargs = dict(extract_cfg=cfg.to_dict(), ds_name=ds_name, split_type=split_type, f=f, date=pd.Timestamp.now().isoformat(),)
-    
-    torch_cols = ['input_ids', 'attention_mask', 'choice_ids', 'question', 'answer_choices', 'example_i', 'label_true', 'sys_instr_name', 'template_name', 'instructed_to_lie']
+    dataset_name = sanitize_filename(
+        f"{cfg.model}_{ds_name}_{split_type}_{N}", replacement_text="_"
+    )
+    f = str(root_folder / ".ds" / f"{dataset_name}")
+    logger.info(
+        f"Creating dataset {dataset_name} with {len(ds_tokens)} examples at `{f}`"
+    )
+
+    info_kwargs = dict(
+        extract_cfg=cfg.to_dict(),
+        ds_name=ds_name,
+        split_type=split_type,
+        f=f,
+        date=pd.Timestamp.now().isoformat(),
+    )
+
+    torch_cols = [
+        "input_ids",
+        "attention_mask",
+        "choice_ids",
+        "question",
+        "answer_choices",
+        "example_i",
+        "label_true",
+        "sys_instr_name",
+        "template_name",
+        "instructed_to_lie",
+    ]
     ds_t_subset = ds_keep_cols(ds_tokens, torch_cols)
     ds = ds_t_subset.to_iterable_dataset()
     # pipeline_it = rep_control_pipeline2(ds, batch_size=batch_size, **text_gen_kwargs)
-    
+
     # first we make the calibration dataset with no intervention
     gen_kwargs = dict(
         model_inputs=ds,
         activations=activations,
         batch_size=batch_size,
     )
-    
+
     if debug:
         # this allow us to debug in a single thread
         pipeline(**gen_kwargs)
-    
+
     dataset_features = get_features(cfg, model.config)
     ds1 = datasets.Dataset.from_generator(
         generator=pipeline,
@@ -216,36 +283,60 @@ def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch
 
 if cfg.disable_ds_cache:
     from datasets import disable_caching
+
     disable_caching()
-    
+
 from src.datasets.load import ds2df, load_ds, get_ds_name, filter_ds_to_known, qc_ds
 
-activations=[activations1, activations2]
+activations = [activations1, activations2]
 for ds_name in cfg.datasets:
-    
     # load dataset
-    N=sum(cfg.max_examples)
-    ds_tokens = load_preproc_dataset(ds_name, tokenizer, N=N, seed=cfg.seed, num_shots=cfg.num_shots, max_length=cfg.max_length, prompt_format=cfg.prompt_format)
+    N = sum(cfg.max_examples)
+    ds_tokens = load_preproc_dataset(
+        ds_name,
+        tokenizer,
+        N=N,
+        seed=cfg.seed,
+        num_shots=cfg.num_shots,
+        max_length=cfg.max_length,
+        prompt_format=cfg.prompt_format,
+    )
+
+    N_train_split = (len(ds_tokens) - N_fit_examples) // 2
 
-    N_train_split = (len(ds_tokens) - N_fit_examples) //2
-    
     N_train_split = cfg.max_examples[0]
 
     # split the dataset, it's preshuffled
     dataset_fit = ds_tokens.select(range(N_fit_examples))
     dataset_train = ds_tokens.select(range(N_fit_examples, N_train_split))
     dataset_test = ds_tokens.select(range(N_train_split, len(ds_tokens)))
-    assert len(dataset_train)>3, f"dataset_train is too small {len(dataset_train)}"
-    assert len(dataset_test)>3
-    
+    assert len(dataset_train) > 3, f"dataset_train is too small {len(dataset_train)}"
+    assert len(dataset_test) > 3
+
     # FIXME:
     # test_intervention_quality(dataset_train)
 
-    ds1, f = create_hs_ds(ds_name, dataset_train, rep_control_pipeline2, split_type="train", debug=True, batch_size=batch_size, activations=activations)
+    ds1, f = create_hs_ds(
+        ds_name,
+        dataset_train,
+        rep_control_pipeline2,
+        split_type="train",
+        debug=True,
+        batch_size=batch_size,
+        activations=activations,
+    )
     clear_mem()
-    ds1, f = create_hs_ds(ds_name, dataset_test, rep_control_pipeline2, split_type="test", debug=True, batch_size=batch_size, activations=activations)
+    ds1, f = create_hs_ds(
+        ds_name,
+        dataset_test,
+        rep_control_pipeline2,
+        split_type="test",
+        debug=True,
+        batch_size=batch_size,
+        activations=activations,
+    )
     clear_mem()
-    
+
     try:
         qc_ds(ds1)
     except:
@@ -258,9 +349,6 @@ for ds_name in cfg.datasets:
 
 # %% [markdown]
 # # To Datasets
-# 
+#
 
 # %%
-
-
-

pyproject.toml

@@ -29,6 +29,8 @@ pytorch-optimizer = "^2.12.0"
 pathvalidate = "^3.2.0"
 torchinfo = "^1.8.0"
 jaxtyping = "^0.2.24"
+autoawq = "^0.1.7"
+bitsandbytes = "^0.41.3.post2"
 
 [[tool.poetry.source]]
 name = "pytorch"

src/datasets/dm.py

@@ -67,6 +67,7 @@ class imdbHSDataModule(pl.LightningDataModule):
             hs = hs.diff(1, axis=1) # this makes it the residual between layers
         if self.skip_layers:
             hs = hs[:, self.skip_layers:] # drop the first 10 layers to prevent overfitting?
+        del hs
         self.hs0 = hs[..., 0]
         self.hs1 = hs[..., 1]
         

src/datasets/intervene.py

@@ -57,7 +57,7 @@ def create_cache_interventions(model, tokenizer, cfg, N_fit_examples=20, batch_s
     tokenizer_args=dict(padding="max_length", max_length=cfg.max_length, truncation=True, add_special_tokens=True)
     
     model_name = cfg.model.replace('/', '-')
-    intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}_{"-" if get_negative else "+"}.pkl'
+    intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}_{"-" if get_negative else "+"}_{direction_method}.pkl'
     intervention_f.parent.mkdir(exist_ok=True, parents=True)
     if not intervention_f.exists():        
         

src/datasets/scores.py

@@ -48,8 +48,8 @@ def choice2id(tokenizer, c: str, whitespace_first=False) -> List[int]:
     
     # Note some tokenizers differentiate between "yes", "\nyes" and " yes", and ideally we want all! 
     ids2 = [
-        tokenizer(f' {c}', add_special_tokens=False)["input_ids"][1],
-        tokenizer(f'\n{c}', add_special_tokens=False)["input_ids"][2],
+        tokenizer(f' {c}', add_special_tokens=False)["input_ids"][-1],
+        tokenizer(f'\n{c}', add_special_tokens=False)["input_ids"][-1],
         tokenizer(f'{c}', add_special_tokens=False)["input_ids"][0],
     ]
     ids = list(set(ids2))

src/extraction/config.py

@@ -12,12 +12,18 @@ class ExtractConfig(Serializable):
     # model: str = "TheBloke/WizardCoder-Python-13B-V1.0-GPTQ"
     # model: str = "TheBloke/Wizard-Vicuna-13B-Uncensored-GPTQ"
     # model: str = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPTQ"
-    model: str = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ" # it wont lie? wtf
+    # model: str = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ" # it wont lie? wtf
+    # model: str = "microsoft/phi-2"
+    # model: str = "microsoft/phi-2"
+    model: str = "/media/wassname/SGIronWolf/projects5/elk/phi-2"
+    
     # model: str = "TheBloke/Llama-2-13B-chat-GPTQ"
     """HF model string identifying the language model to extract hidden states from."""
 
-    prompt_format: str | None = None
-    """llama, llama2, chatml, as a backup to tokenizer see structure.yaml file."""
+    batch_size: int = 6
+
+    prompt_format: str | None = 'phi'
+    """if the tokenizer does not have a chat template you can set a custom one. see src/prompts/templates/prompt_formats/readme.md."""
     
     data_dirs: tuple[str, ...] = ()
     """Directory to use for caching the hiddens. Defaults to `HF_DATASETS_CACHE`."""
@@ -53,3 +59,11 @@ class ExtractConfig(Serializable):
     
     disable_ds_cache: bool = False
     """Disable huggingface datasets cache."""
+
+    intervention_direction_method: str = "cluster_mean"
+    """"how to intervent: pca, cluster_mean, random"""
+
+    intervention_fit_examples: int = 60
+    """how many example to use for intervention calibration"""
+
+

src/models/load.py

@@ -10,6 +10,7 @@ import torch
 from src.datasets.dropout import check_for_dropout
 from loguru import logger
 from typing import Tuple
+from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 
 def verbose_change_param(tokenizer, path, after):
     
@@ -20,7 +21,7 @@ def verbose_change_param(tokenizer, path, after):
     return tokenizer
 
 
-def load_model(model_repo = "TheBloke/WizardCoder-Python-13B-V1.0-GPTQ") -> Tuple[AutoModelForCausalLM, PreTrainedTokenizerBase]:
+def load_model(model_repo =  "microsoft/phi-2") -> Tuple[AutoModelForCausalLM, PreTrainedTokenizerBase]:
     """
     A uncensored and large coding ones might be best for lying.
     
@@ -30,10 +31,12 @@ def load_model(model_repo = "TheBloke/WizardCoder-Python-13B-V1.0-GPTQ") -> Tupl
     model_options = dict(
         device_map="auto",
         torch_dtype=torch.float16, 
+        # load_in_8bit=True,
+        trust_remote_code=True,
     )
 
-    config = AutoConfig.from_pretrained(model_repo, use_cache=False)
-    verbose_change_param(config, 'use_cache', False)
+    config = AutoConfig.from_pretrained(model_repo, trust_remote_code=True,)
+    # verbose_change_param(config, 'use_cache', False)
     
     tokenizer = AutoTokenizer.from_pretrained(model_repo, use_fast=True, legacy=False)
     verbose_change_param(tokenizer, 'pad_token_id', 0)
@@ -41,6 +44,7 @@ def load_model(model_repo = "TheBloke/WizardCoder-Python-13B-V1.0-GPTQ") -> Tupl
     verbose_change_param(tokenizer, 'truncation_side', 'left')
     
     model = AutoModelForCausalLM.from_pretrained(model_repo, config=config, 
+        # gptq_config=gptq_config,
                                                  **model_options)
 
     return model, tokenizer

src/prompts/prompt_loading.py

@@ -35,6 +35,8 @@ TEMPLATES_FOLDER_PATH = Path(__file__).parent / "templates"
 def load_prompt_structure(prompt_format='llama2'):
     # for use with https://huggingface.co/docs/transformers/main/chat_templating is the tokenizer doesn't include it
     f = TEMPLATES_FOLDER_PATH  / "prompt_formats" / f"{prompt_format}.jinja2"
+    if not f.exists():
+        raise FileNotFoundError(f"Could not find prompt format {prompt_format} at {f}")
     return f.open().read()
 
 
@@ -296,7 +298,7 @@ def _convert_to_prompts(
 
 
 
-def load_preproc_dataset(ds_name: str, tokenizer: PreTrainedTokenizerBase, N:int, prompt_format:str, split_type:str="train", seed=42, num_shots=1, max_length=999) -> Dataset:
+def load_preproc_dataset(ds_name: str, tokenizer: PreTrainedTokenizerBase, N:int, prompt_format:str = None, split_type:str="train", seed=42, num_shots=1, max_length=999) -> Dataset:
     """load a preprocessed dataset of tokens."""
     ds_prompts = Dataset.from_generator(
         load_prompts,
@@ -312,7 +314,7 @@ def load_preproc_dataset(ds_name: str, tokenizer: PreTrainedTokenizerBase, N:int
     )
     
     
-    if tokenizer.chat_template is None:
+    if prompt_format:
         tokenizer.chat_template = load_prompt_structure(prompt_format=prompt_format)
 
     # ## Format prompts

src/prompts/templates/prompt_formats/phi.jinja2

@@ -0,0 +1,2 @@
+{# https://huggingface.co/microsoft/phi-2 #}
+{% for message in messages %}{% if message['role'] == 'user' %}{{ 'Instruct: ' + message['content'] + ' Output: ' }}{% elif message['role'] == 'assistant'%}{{ message['content'] + ' ' }}{% elif message['role'] == 'system'%}{{ 'System: ' + message['content'] + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!')}}{% endif %}{% endfor %}

src/prompts/templates/prompt_formats/vicuna.jinja2

@@ -1,3 +1,3 @@
 {# https://huggingface.co/TheBloke/Wizard-Vicuna-13B-Uncensored-GPTQ #}
-{# not some vicuna variant are diff #}
+{# note some vicuna variant are diff #}
 {% for message in messages %}{% if message['role'] == 'user' %}{{ 'USER: ' + message['content'] + ' ASSISTANT: ' }}{% elif message['role'] == 'assistant'%}{{ message['content'] + ' ' }}{% elif message['role'] == 'system'%}{{ 'SYSTEM: ' + message['content'] + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!')}}{% endif %}{% endfor %}

src/repe/rep_readers.py

@@ -71,15 +71,15 @@ class RepReader(ABC):
 
         if self.needs_hiddens and hidden_states is not None and len(hidden_states) > 0:
             for layer in hidden_layers:    
-                assert hidden_states[layer].shape[0] == 2 * len(train_choices), f"Shape mismatch between hidden states ({hidden_states[layer].shape[0]}) and labels ({len(train_choices)})"
+                assert hidden_states[layer].shape[0] == len(train_choices), f"Shape mismatch between hidden states ({hidden_states[layer].shape[0]}) and labels ({len(train_choices)})"
                 
                 signs[layer] = []
                 for component_index in range(self.n_components):
                     transformed_hidden_states = project_onto_direction(hidden_states[layer], self.directions[layer][component_index])
                     projected_scores = [transformed_hidden_states[i:i+2] for i in range(0, len(transformed_hidden_states), 2)]
 
-                    outputs_min = [1 if min(o) == o[label] else 0 for o, label in zip(projected_scores, train_choices)]
-                    outputs_max = [1 if max(o) == o[label] else 0 for o, label in zip(projected_scores, train_choices)]
+                    outputs_min = [1 if min(o) == o[int(label)] else 0 for o, label in zip(projected_scores, train_choices)]
+                    outputs_max = [1 if max(o) == o[int(label)] else 0 for o, label in zip(projected_scores, train_choices)]
                     
                     signs[layer].append(-1 if np.mean(outputs_min) > np.mean(outputs_max) else 1)
         else:
@@ -189,6 +189,8 @@ class ClusterMeanRepReader(RepReader):
 
     def get_rep_directions(self, model, tokenizer, hidden_states, hidden_layers, **kwargs):
 
+        # see also https://github.com/likenneth/honest_llama/blob/207bb14b2c005e0593487cca8d22e072cbcb987b/utils.py#L730
+
         # train labels is necessary to differentiate between different classes
         train_choices = kwargs['train_choices'] if 'train_choices' in kwargs else None
         assert train_choices is not None, "ClusterMeanRepReader requires train_choices to differentiate two clusters"

src/repe/rep_reading_pipeline.py

@@ -138,10 +138,10 @@ class RepReadingPipeline(Pipeline):
             
             # get differences between pairs
             relative_hidden_states = {k: np.copy(v) for k, v in hidden_states.items()}
-            for layer in hidden_layers:
+            for layer in hidden_layers[1:]:
                 for _ in range(n_difference):
-                    # TODO: check this, it's even - odd? on what dimension?
-                    relative_hidden_states[layer] = relative_hidden_states[layer][::2] - relative_hidden_states[layer][1::2]
+                    # FIXME: this is wrong, it's skipping batches...
+                    relative_hidden_states[layer] = relative_hidden_states[layer] - relative_hidden_states[layer-1]
 
 		# get the directions
         direction_finder.directions = direction_finder.get_rep_directions(