Spaces:
Sleeping
Sleeping
Update LLMPipeline.py
Browse files- LLMPipeline.py +2 -0
LLMPipeline.py
CHANGED
|
@@ -103,11 +103,13 @@ model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
|
|
| 103 |
# print(tokenizer.decode(outputs[0]))
|
| 104 |
|
| 105 |
def generate_image_prompt(summary):
|
|
|
|
| 106 |
system_prompt_rewrite = "You are a creative meme prompt generator. Create a visual prompt to generate a meme of a message. Do not return any text other than the new prompt."
|
| 107 |
user_prompt_rewrite = "Create a visual meme prompt from the message. Keep its main point.\nThe message:"
|
| 108 |
messages = [{"role": "system", "content": system_prompt_rewrite}, {"role": "user", "content":f"{user_prompt_rewrite} {summary}"}]
|
| 109 |
input_text=tokenizer.apply_chat_template(messages, tokenize=False)
|
| 110 |
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
|
| 111 |
outputs = model.generate(inputs, max_new_tokens=500, temperature=0.2, top_p=0.9, do_sample=True)
|
|
|
|
| 112 |
return tokenizer.decode(outputs[0])
|
| 113 |
|
|
|
|
| 103 |
# print(tokenizer.decode(outputs[0]))
|
| 104 |
|
| 105 |
def generate_image_prompt(summary):
|
| 106 |
+
print('SUMMARY:', summary)
|
| 107 |
system_prompt_rewrite = "You are a creative meme prompt generator. Create a visual prompt to generate a meme of a message. Do not return any text other than the new prompt."
|
| 108 |
user_prompt_rewrite = "Create a visual meme prompt from the message. Keep its main point.\nThe message:"
|
| 109 |
messages = [{"role": "system", "content": system_prompt_rewrite}, {"role": "user", "content":f"{user_prompt_rewrite} {summary}"}]
|
| 110 |
input_text=tokenizer.apply_chat_template(messages, tokenize=False)
|
| 111 |
inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
|
| 112 |
outputs = model.generate(inputs, max_new_tokens=500, temperature=0.2, top_p=0.9, do_sample=True)
|
| 113 |
+
print('RESUMIDO')
|
| 114 |
return tokenizer.decode(outputs[0])
|
| 115 |
|