initial commit

.gitignore

@@ -0,0 +1 @@
+.ipynb_checkpoints/

app.py

@@ -37,12 +37,10 @@ if not torch.cuda.is_available():
 if torch.cuda.is_available():
     model_id = "meta-llama/Llama-2-7b-chat-hf" # not gated version.
     model = AutoModelForCausalLM.from_pretrained(
-        model_id, device_map="auto", torch_dtype=torch.bfloat16
+        model_id, device_map="cuda", torch_dtype=torch.bfloat16
     )
     reft_model = ReftModel.load("pyvene/reft_goody2", model, from_huggingface_hub=True)
-    # a little hacky.
-    for k, v in reft_model.interventions.items():
-        v[0].to(model.device)
+    reft_model.set_device("cuda")
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.use_default_system_prompt = True
 
@@ -77,7 +75,7 @@ def generate(
     
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = {
-        "base": {"input_ids": prompt["input_ids"], "attention_mask": prompt["attention_mask"]},
+        "base": {"input_ids": input_ids, "attention_mask": attention_mask},
         "unit_locations": {"sources->base": (None, [[[base_unit_location]]])},
         "max_new_tokens": max_new_tokens,
         "intervene_on_prompt": True,