Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,13 +7,12 @@ st.set_page_config(page_title="FitPlan AI", layout="centered")
 # LOAD MODEL 
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
-    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
+    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-
 # TITLE
 st.title(" FitPlan AI – User Fitness Profile")
 
@@ -130,70 +129,64 @@ fitness_level = st.radio(
     horizontal=True
 )
 
-def create_prompt(gender, bmi, bmi_status, goal, level, equipment):
-
-    equipment_list = ", ".join(equipment)
-
-    prompt = f"""
-Instruction:
-You are a certified professional fitness trainer.
-
-Generate a structured 5-day workout plan.
-
-User Information:
-Gender: {gender}
-BMI: {bmi:.2f} ({bmi_status})
-Goal: {goal}
-Fitness Level: {level}
-Equipment Available: {equipment_list}
-
-Rules:
-1. Only generate Day 1 to Day 5.
-2. Each day must include:
-   - Focus
-   - 4 exercises
-   - Sets x Reps
-   - Rest time
-3. Do not repeat days.
-4. Do not generate explanations.
-5. Keep output clean and professional.
-
-Start directly with:
-Day 1:
-"""
+# SUBMIT BUTTON
+if st.button(" Submit Profile"):
 
-    return prompt
+    if not name:
+        st.error("Please enter your name.")
 
-def generate_workout_plan(prompt):
+    elif height <= 0 or weight <= 0:
+        st.error("Please enter valid height and weight.")
 
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    elif not equipment:
+        st.error("Please select at least one equipment option.")
 
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=400,
-        num_beams=4,
-        early_stopping=True,
-        no_repeat_ngram_size=3
-    )
+    else:
+        st.success("✅ Profile Submitted Successfully!")
 
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        bmi_status = bmi_category(bmi)
+        equipment_list = ", ".join(equipment)
 
-    return result
+        prompt = f"""
+You are a certified professional fitness trainer.
+Generate a structured 5-day workout plan based on the following user profile.
+User Profile:
+- Gender: {gender}
+- BMI: {bmi:.2f} ({bmi_status})
+- Goal: {goal}
+- Fitness Level: {fitness_level}
+- Available Equipment: {equipment_list}
+Instructions:
+- Divide into Day 1 to Day 5
+- Use EXACT format below
+- Do not change format
+- Do not add extra text
+Strict Format:
+Day 1: <Workout Focus>
+1. <Exercise Name>
+   Sets: <number>
+   Reps: <range>
+   Rest: <seconds>
+2. <Exercise Name>
+   Sets: <number>
+   Reps: <range>
+   Rest: <seconds>
+Repeat same structure for all 5 days.
+Only return the workout plan.
+"""
 
+        with st.spinner("Generating your AI workout plan..."):
 
-if st.button("Submit Profile"):
+            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
-    prompt = create_prompt(
-        gender,
-        bmi,
-        bmi_status,
-        goal,
-        fitness_level,
-        equipment
-    )
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=300,
+                temperature=0.7,
+                do_sample=True
+            )
 
-    with st.spinner("Generating AI workout plan..."):
-        result = generate_workout_plan(prompt)
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
-    st.subheader("Your Personalized Workout Plan")
-    st.write(result)
+        st.subheader("🏋️ Your Personalized Workout Plan")
+        st.write(result)