Tonic commited on
Commit
58e5994
1 Parent(s): 6cbc4c5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -21,8 +21,8 @@ model = AutoModelForCausalLM.from_pretrained(model_path, device_map='cuda', quan
21
  def generate_text(prompt, custom_prompt, temperature, max_length):
22
  structured_input = f"[[[Title]]]: {prompt}\n[[[Content]]]: {custom_prompt}"
23
  inputs = tokenizer(structured_input, return_tensors='pt').to('cuda')
24
- generated_ids = model.generate(**inputs, max_new_tokens=max_length, temperature=temperature, do_sample=True).split(tokenizer.eos_token)[0]
25
- generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=False)
26
  return generated_text
27
 
28
 
 
21
  def generate_text(prompt, custom_prompt, temperature, max_length):
22
  structured_input = f"[[[Title]]]: {prompt}\n[[[Content]]]: {custom_prompt}"
23
  inputs = tokenizer(structured_input, return_tensors='pt').to('cuda')
24
+ generated_ids = model.generate(**inputs, max_new_tokens=max_length, temperature=temperature, do_sample=True)
25
+ generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=False).split(tokenizer.eos_token)[0])
26
  return generated_text
27
 
28