Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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 =
|
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 =
|
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
|