Update app.py

app.py

@@ -66,8 +66,8 @@ def get_hidden_states(raw_original_prompt, progress=gr.Progress()):
 
 
 @spaces.GPU
-def generate_interpretation_gpu(interpret_prompt, **kwargs):
-    return interpret_prompt.generate(**kwargs)
+def generate_interpretation_gpu(interpret_prompt, *args, **kwargs):
+    return interpret_prompt.generate(*args, **kwargs)
 
 
 def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
@@ -94,7 +94,8 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
 
     # generate the interpretations
-    generate = generate_interpretation_gpu if use_gpu else lambda interpretation_prompt, **kwargs: interpretation_prompt.generate(**kwargs)
+    generate = generate_interpretation_gpu if use_gpu else lambda interpretation_prompt, *args, 
+                                                                    **kwargs: interpretation_prompt.generate(*args, **kwargs)
     generated = generate(interpretation_prompt, model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''