Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,11 +1,12 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+
 # -------------------------
 # PAGE CONFIG
-
 # -------------------------
 st.set_page_config(page_title="FitPlan-AI", page_icon="💪")
+
 # -------------------------
 # BMI FUNCTIONS
 # -------------------------
@@ -31,7 +32,7 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
     model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
     return tokenizer, model
- 
+
 tokenizer, model = load_model()
 
 # -------------------------
@@ -42,6 +43,7 @@ st.title("💪 FitPlan-AI: Personalized Fitness Profile")
 with st.form("fitness_form"):
 
     st.subheader("Personal Information")
+
     name = st.text_input("Full Name*", placeholder="Enter your name")
 
     col1, col2 = st.columns(2)
@@ -68,15 +70,10 @@ with st.form("fitness_form"):
     )
 
     submit = st.form_submit_button("Submit Profile")
-# -------------------------
-# HANDLE SUBMISSION
-# -------------------------
-# SUBMIT BUTTON
 
 # -------------------------
 # HANDLE SUBMISSION
 # -------------------------
-
 if submit:
 
     if not name:
@@ -91,23 +88,36 @@ if submit:
     else:
         st.success("Profile Submitted Successfully!")
 
-        # ✅ Calculate BMI
+        # Calculate BMI
         bmi = calculate_bmi(weight, height)
         bmi_status = get_category(bmi)
 
+        st.write(f"**Your BMI:** {bmi} ({bmi_status})")
+
         equipment_list = ", ".join(equipment)
 
+        # Strong structured prompt
         prompt = f"""
 You are a certified professional fitness trainer.
 
-Create a detailed 5-day workout plan.
+Create a COMPLETE detailed 5-day workout plan.
+
+IMPORTANT:
+- You MUST generate Day 1, Day 2, Day 3, Day 4, and Day 5.
+- Each day must include:
+  - Warm-up
+  - 4-6 exercises
+  - Sets and reps
+  - Cool down
+- Do NOT stop after Day 1.
+- Write all 5 days clearly.
 
 User Information:
-- Name: {name}
-- BMI: {bmi} ({bmi_status})
-- Goal: {goal}
-- Fitness Level: {level}
-- Equipment Available: {equipment_list}
+Name: {name}
+BMI: {bmi} ({bmi_status})
+Goal: {goal}
+Fitness Level: {level}
+Equipment Available: {equipment_list}
 
 Start directly with:
 
@@ -120,12 +130,13 @@ Day 1:
 
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=600,
-                temperature=0.7,
-                do_sample=True
+                max_new_tokens=1200,
+                temperature=0.8,
+                do_sample=True,
+                repetition_penalty=1.1
             )
 
             result = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
-        st.subheader("Your Personalized Workout Plan")
+        st.subheader("🏋️ Your Personalized 5-Day Workout Plan")
         st.write(result)