Update app.py

app.py

@@ -38,12 +38,13 @@ Assistant: "girlfriend;mother;father;friend"
 
 @spaces.GPU
 def get_past_key_values(system_prompt):
+    model, tokenizer = pipe.model, pipe.tokenizer
     messages = [{'role': 'system', 'content': system_prompt}]
-    tokenized = pipe.tokenizer.apply_chat_template(messages, return_tensors='pt')
+    tokenized = tokenizer.apply_chat_template(messages, return_tensors='pt')
     
     # assert that this is indeed a prefix (TODO: make sure this is robust)
     messages.append({'role': 'user', 'content': 'TEST'})
-    tokenized_test = pipe.tokenizer.apply_chat_template(messages, return_tensors='pt')
+    tokenized_test = tokenizer.apply_chat_template(messages, return_tensors='pt')
     assert (tokenized_test[:, :tokenized.shape[1]] == tokenized).all().cpu().item()
     
     return model(**tokenized.to(model.device)).past_key_values