Update app.py

app.py

@@ -18,19 +18,23 @@ quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 model = AutoModelForCausalLM.from_pretrained(model_path, device_map='cuda', quantization_config=quantization_config)
 
 @spaces.GPU
-def generate_text(prompt, custom_prompt, temperature, max_length):
-    structured_input = f"[[[Title]]]: {prompt}\n[[[Content]]]: {custom_prompt}"
-    inputs = tokenizer(structured_input, return_tensors='pt').to('cuda')
-    generated_ids = model.generate(**inputs, max_new_tokens=max_length, temperature=temperature, do_sample=True)
-    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=False) #.split(tokenizer.eos_token)[0]
+def generate_text(usertitle, content, max_new_tokens=512,model=model, tokenizer=tokenizer, temperature=0.7):
+    msg = [{
+        'title': title,
+        'content': content
+    }]
+    inputs = tokenizer.apply_chat_template((msg, return_tensors='pt').cuda())    
+    generated_ids = model.generate(inputs['input_ids'], max_new_tokens=max_new_tokens, temperature=temperature, pad_token_id=tokenizer.eos_token_id)
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    
     return generated_text
 
 
 def gradio_app():
     with gr.Blocks() as demo:
         gr.Markdown(title)
-        custom_prompt = gr.Textbox(label="WordPhrases", value=examplecofee, lines=5)
-        prompt = gr.Textbox(label="Title", value="Cortado", lines=1)
+        usertitle = gr.Textbox(label="Title", value="Cortado", lines=1)
+        content = gr.Textbox(label="WordPhrases", value=examplecofee, lines=5)
         with gr.Row():
             temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
             max_length = gr.Slider(minimum=250, maximum=1024, step=10, value=450, label="Generate Length")
@@ -39,7 +43,7 @@ def gradio_app():
 
         generate_btn.click(
             fn=generate_text,
-            inputs=[prompt, custom_prompt, temperature, max_length],
+            inputs=[usertitle, content, temperature, max_length],
             outputs=output
         )