Spaces:

pyvene
/

AxBench-ReFT-r1-16K

Running on Zero

frankaging commited on Apr 23, 2024

Commit

327242b

1 Parent(s): a4bd2e7

initial commit

Files changed (1) hide show

app.py CHANGED Viewed

@@ -85,7 +85,7 @@ def generate(
     prompt = tokenizer.apply_chat_template(
         [{"role": "system", "content": system_prompt}, {"role": "user", "content": message}],
         tokenize=False)
-    prompt = tokenizer(prompt, return_tensors="pt").to(device)
     unit_locations = torch.IntTensor([pyreft.get_intervention_locations(
         last_position=prompt["input_ids"].shape[-1],

     prompt = tokenizer.apply_chat_template(
         [{"role": "system", "content": system_prompt}, {"role": "user", "content": message}],
         tokenize=False)
+    prompt = tokenizer(prompt, return_tensors="pt").to(model.device)
     unit_locations = torch.IntTensor([pyreft.get_intervention_locations(
         last_position=prompt["input_ids"].shape[-1],