Tonic commited on
Commit
6a5a0e9
Β·
verified Β·
1 Parent(s): ddecc12

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -9
app.py CHANGED
@@ -23,15 +23,9 @@ model = AutoModelForCausalLM.from_pretrained( model_path, device_map='cuda', loa
23
 
24
  @spaces.GPU
25
  def generate_text(prompt, custom_prompt, temperature, max_length):
26
- # Example message format
27
- msg = [{
28
- 'title': prompt,
29
- 'content': custom_prompt
30
- }]
31
-
32
- prompt_text = f"title: {msg[0]['title']}\ncontent: {msg[0]['content']}"
33
- inputs = tokenizer(prompt_text, return_tensors='pt').to('cuda')
34
- generated_ids = model.generate(**inputs, max_new_tokens=max_length, temperature=temperature, pad_token_id=tokenizer.eos_token_id)
35
  generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
36
 
37
  return generated_text
 
23
 
24
  @spaces.GPU
25
  def generate_text(prompt, custom_prompt, temperature, max_length):
26
+ structured_input = f"[Title]: {prompt}\n[Content]: {custom_prompt}"
27
+ inputs = tokenizer(structured_input, return_tensors='pt').to('cuda')
28
+ generated_ids = model.generate(**inputs, max_new_tokens=max_length, temperature=temperature, do_sample=True, pad_token_id=tokenizer.eos_token_id)
 
 
 
 
 
 
29
  generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
30
 
31
  return generated_text