Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 from huggingface_hub import InferenceClient
-import os
 
 client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.1")
 
@@ -14,14 +13,6 @@ def format_prompt(message, history, system_prompt=""):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def extract_filename_and_code(bot_response):
-    try:
-        with open("default.py", 'w') as file:
-            file.write(bot_response)
-        st.write("Successfully wrote to default.py!")
-    except Exception as e:
-        st.write(f"Error writing to file: {e}")
-
 def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
     temperature = float(temperature)
     if temperature < 1e-2:
@@ -43,8 +34,7 @@ def generate(prompt, history, system_prompt="", temperature=0.9, max_new_tokens=
 
     for response in stream:
         output += response.token.text
-        extract_filename_and_code(output)
-        yield output
+        st.session_state.generated_text = output
 
     return output
 
@@ -64,12 +54,16 @@ max_new_tokens = st.slider("Max new tokens", 0, 8192, 512, step=64)
 top_p = st.slider("Top-p (nucleus sampling)", 0.0, 1.0, 0.95, step=0.05)
 repetition_penalty = st.slider("Repetition penalty", 1.0, 2.0, 1.0, step=0.05)
 
+# Initialize session state for generated text
+if "generated_text" not in st.session_state:
+    st.session_state.generated_text = ""
+
 # Generate button
 if st.button("Generate Waifu"):
     history = []
     prompt = f"Create a waifu character named {name} with {hair_color} hair, a {personality} personality, and wearing a {outfit_style}."
-    output = generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+    generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
     
     # Display the generated character
     st.subheader("Generated Waifu Character")
-    st.write(next(output))
+    st.write(st.session_state.generated_text)