Update app.py

app.py

@@ -76,12 +76,10 @@ def user_input(user_question):
     # Extract the thought process
     thought_process = ""
     if "<think>" in response['output_text'] and "</think>" in response['output_text']:
-        thought_process_match = re.search(r"<think>(.*?)</think>", response['output_text'], re.DOTALL)
-        if thought_process_match:
-            thought_process = thought_process_match.group(1).strip()
+        thought_process = response['output_text'].split("<think>")[1].split("</think>")[0].strip()
 
     # Remove the thought process from the main response
-    clean_response = re.sub(r"<think>.*?</think>", "", response['output_text']).strip()
+    clean_response = response['output_text'].replace(f"<think>{thought_process}</think>", "").strip()
 
     # Display the model's thought process in the expander
     with st.expander("Model Thought Process"):
@@ -128,3 +126,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+