Update app.py

app.py

@@ -65,6 +65,11 @@ def get_hidden_states(raw_original_prompt, progress=gr.Progress()):
     return [progress_dummy_output, hidden_states, *token_btns]
 
 
+@spaces.GPU
+def generate_interpretation_gpu(interpret_prompt, **kwargs):
+    return interpret_prompt.generate(**kwargs)
+
+
 def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens, do_sample, 
                        temperature, top_k, top_p, repetition_penalty, length_penalty, use_gpu, i, 
                        num_beams=1):
@@ -89,7 +94,7 @@ def run_interpretation(global_state, raw_interpretation_prompt, max_new_tokens,
     interpretation_prompt = InterpretationPrompt(tokenizer, interpretation_prompt)
 
     # generate the interpretations
-    generate = spaces.GPU(interpretation_prompt.generate) if use_gpu else interpretation_prompt.generate
+    generate = generate_interpretation_gpu if use_gpu else lambda lambda interpretation_prompt, **kwargs: interpretation_prompt.generate(**kwargs)
     generated = generate(model, {0: interpreted_vectors}, k=3, **generation_kwargs)
     generation_texts = tokenizer.batch_decode(generated)
     progress_dummy_output = ''