Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Optional, Tuple, Any
 from functools import partial
 import gradio as gr
 import spaces
@@ -53,7 +53,7 @@ start_messages = [
 # functions
 @dataclass
 class PastKV:
-    past_key_values: Optional[torch.Tensor] = None
+    past_key_values: Any = None
 
 past_key_values = PastKV()
 
@@ -82,7 +82,7 @@ def generate(text, past_key_values):
         {'role': 'user', 'content': text}
     ]
     response = pipe(messages, 
-                    past_key_values=past_key_values, #past_kv_to_device(past_key_values, pipe.model.device), 
+                    past_key_values=past_key_values.past_key_values, #past_kv_to_device(past_key_values, pipe.model.device), 
                     **generate_kwargs)[0]['generated_text']
     return response[-1]['content']
 
@@ -90,6 +90,7 @@ def generate(text, past_key_values):
 if __name__ == "__main__":
     with torch.no_grad():
         set_past_key_values()
-        demo = gr.Interface(partial(generate, past_key_values=past_key_values.past_key_values), 
+        print(past_key_values)
+        demo = gr.Interface(partial(generate), 
                             inputs="textbox", outputs="textbox")
         demo.launch()