Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import subprocess
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import re
 
 st.title('Eleanor Rigby')
 
@@ -22,9 +23,10 @@ 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]
-    formatted_output = output.split("Output:")
-    trimmed_output = formatted_output[1].strip()
-    st.text("Generated Result:")
-    st.write(trimmed_output)
+    match = re.search(r'Output:(.*)', output)
+    if match:
+        output_content = match.group(1).strip()
+        st.text("Generated Result:")
+        st.write(output_content)
     
 st.button('Generate Result', on_click=run_inference)