Update app.py

app.py

@@ -15,7 +15,7 @@ user_prompt = st.text_area("Enter your prompt that can be song lyrics e.g. 'Yest
 
 def run_inference():
     instruct_prompt = "Instruct:You are a song writer and your main reference is The Beatles. Write a song lyrics by completing these words:"
-    output_prompt = "Output:\n"
+    output_prompt = "Output:"
     input = inference_tokenizer(""" {0}{1}\n{2} """.format(instruct_prompt, user_prompt, output_prompt),
                      return_tensors="pt",
                      return_attention_mask=False,
@@ -23,10 +23,7 @@ def run_inference():
                      truncation=True)
     result = inference_model.generate(**input, repetition_penalty=1.2, max_length=1024)
     output = inference_tokenizer.batch_decode(result, skip_special_tokens=True)[0]
-    match = re.search(r'Output:(.*)', output)
-    if match:
-        output_content = match.group(1).strip()
-        st.text("Generated Result:")
-        st.write(output_content)
+    st.text("Generated Result:")
+    st.write(output)
     
 st.button('Generate Result', on_click=run_inference)