dar-tau commited on
Commit
3824c46
·
verified ·
1 Parent(s): e0ae9ce

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -38,12 +38,13 @@ Assistant: "girlfriend;mother;father;friend"
38
 
39
  @spaces.GPU
40
  def get_past_key_values(system_prompt):
 
41
  messages = [{'role': 'system', 'content': system_prompt}]
42
- tokenized = pipe.tokenizer.apply_chat_template(messages, return_tensors='pt')
43
 
44
  # assert that this is indeed a prefix (TODO: make sure this is robust)
45
  messages.append({'role': 'user', 'content': 'TEST'})
46
- tokenized_test = pipe.tokenizer.apply_chat_template(messages, return_tensors='pt')
47
  assert (tokenized_test[:, :tokenized.shape[1]] == tokenized).all().cpu().item()
48
 
49
  return model(**tokenized.to(model.device)).past_key_values
 
38
 
39
  @spaces.GPU
40
  def get_past_key_values(system_prompt):
41
+ model, tokenizer = pipe.model, pipe.tokenizer
42
  messages = [{'role': 'system', 'content': system_prompt}]
43
+ tokenized = tokenizer.apply_chat_template(messages, return_tensors='pt')
44
 
45
  # assert that this is indeed a prefix (TODO: make sure this is robust)
46
  messages.append({'role': 'user', 'content': 'TEST'})
47
+ tokenized_test = tokenizer.apply_chat_template(messages, return_tensors='pt')
48
  assert (tokenized_test[:, :tokenized.shape[1]] == tokenized).all().cpu().item()
49
 
50
  return model(**tokenized.to(model.device)).past_key_values