islasher commited on
Commit
5187e93
·
verified ·
1 Parent(s): e0dd62e

Update LLMPipeline.py

Browse files
Files changed (1) hide show
  1. 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