Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,6 +7,8 @@ import torch
 # -------------------------
 st.set_page_config(page_title="FitPlan-AI", page_icon="💪")
 
+st.title("💪 FitPlan-AI: Personalized Fitness Profile")
+
 # -------------------------
 # BMI FUNCTIONS
 # -------------------------
@@ -25,7 +27,7 @@ def get_category(bmi):
         return "Obese"
 
 # -------------------------
-# LOAD FLAN-T5 MODEL
+# LOAD MODEL
 # -------------------------
 @st.cache_resource
 def load_model():
@@ -36,10 +38,8 @@ def load_model():
 tokenizer, model = load_model()
 
 # -------------------------
-# UI STARTS
+# FORM
 # -------------------------
-st.title("💪 FitPlan-AI: Personalized Fitness Profile")
-
 with st.form("fitness_form"):
 
     st.subheader("Personal Information")
@@ -66,7 +66,8 @@ with st.form("fitness_form"):
 
     equipment = st.multiselect(
         "Available Equipment",
-        ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment", "Kettlebell", "Pull-up Bar"]
+        ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment", 
+         "Kettlebell", "Pull-up Bar"]
     )
 
     submit = st.form_submit_button("Submit Profile")
@@ -92,51 +93,60 @@ if submit:
         bmi = calculate_bmi(weight, height)
         bmi_status = get_category(bmi)
 
-        st.write(f"**Your BMI:** {bmi} ({bmi_status})")
+        st.write(f"### 📊 Your BMI: {bmi} ({bmi_status})")
 
         equipment_list = ", ".join(equipment)
 
-        # Strong structured prompt
-        prompt = f"""
+        # -------------------------
+        # GENERATE 5-DAY PLAN (LOOP METHOD)
+        # -------------------------
+        with st.spinner("Generating your 5-day workout plan..."):
+
+            full_plan = ""
+
+            for day in range(1, 6):
+
+                prompt = f"""
 You are a certified professional fitness trainer.
 
-Create a COMPLETE detailed 5-day workout plan.
+Create a detailed workout plan for Day {day}.
 
-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.
+Include:
+- Warm-up
+- 4 to 6 exercises
+- Sets and reps
+- Cool down
 
-User Information:
-Name: {name}
+User Details:
 BMI: {bmi} ({bmi_status})
 Goal: {goal}
 Fitness Level: {level}
 Equipment Available: {equipment_list}
 
-Start directly with:
+Start with:
 
-Day 1:
+Day {day}:
 """
 
-        with st.spinner("Generating your AI workout plan..."):
+                inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
-            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=350,
+                    temperature=0.8,
+                    do_sample=True,
+                    repetition_penalty=1.1
+                )
 
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=1200,
-                temperature=0.8,
-                do_sample=True,
-                repetition_penalty=1.1
-            )
+                result = tokenizer.decode(
+                    outputs[0], 
+                    skip_special_tokens=True
+                ).strip()
 
-            result = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+                full_plan += result + "\n\n"
 
+        # -------------------------
+        # DISPLAY PLAN
+        # -------------------------
         st.subheader("🏋️ Your Personalized 5-Day Workout Plan")
-        st.write(result)
+        st.write(full_plan)