Update app.py

app.py

@@ -1,8 +1,10 @@
 import streamlit as st
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-# Streamlit UI
+
+# Must be the first Streamlit command
 st.set_page_config(page_title="Email Writer with FLAN-T5", layout="centered")
-# Load model and tokenizer
+
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+
 @st.cache_resource
 def load_model():
     model_name = "google/flan-t5-large"
@@ -12,27 +14,30 @@ def load_model():
 
 tokenizer, model = load_model()
 
-
+# Streamlit UI
 st.title("📧 AI Email Writer")
 st.write("Generate professional emails with the help of `google/flan-t5-large` model.")
 
-# Input box
+# Input section
 user_instruction = st.text_area("✍️ Enter your email instruction", 
                                 placeholder="e.g. Write a polite follow-up email to a client who hasn’t responded in a week.", 
                                 height=150)
 
 temperature = st.slider("🎯 Creativity (Temperature)", 0.1, 1.0, 0.7, 0.1)
-max_tokens = st.slider("📏 Max Tokens", 50, 512, 200, 10)
+max_tokens = st.slider("📏 Max Output Tokens", 50, 512, 200, 10)
 
-# Generate email on button click
 if st.button("Generate Email"):
-    if user_instruction.strip() == "":
-        st.warning("Please enter an instruction.")
+    if not user_instruction.strip():
+        st.warning("Please enter a valid instruction.")
     else:
         with st.spinner("Generating email..."):
-            inputs = tokenizer(user_instruction, return_tensors="pt").input_ids
-            outputs = model.generate(inputs)
-            email_output = tokenizer.decode(outputs[0])
+            input_ids = tokenizer(user_instruction, return_tensors="pt").input_ids
+            output_ids = model.generate(
+                input_ids=input_ids,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                do_sample=True
+            )
+            email_output = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
         st.success("Here’s your generated email:")
         st.text_area("📨 Generated Email", value=email_output, height=200)
-