Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,77 +1,89 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
-
-# -------------------------
-# Page Config FIRST
-# -------------------------
-
-st.set_page_config(page_title="FitPlan AI", page_icon="💪")
-
-# 🎨 Background Style
-st.markdown(
-    """
-    <style>
-    .stApp {
-        background: linear-gradient(to right, #ff512f, #dd2476);
-    }
-    </style>
-    """,
-    unsafe_allow_html=True
-)
-
-# -------------------------
-# Load Model (Small for HF)
-# -------------------------
 
+# -----------------------------
+# PAGE CONFIG
+# -----------------------------
+st.set_page_config(page_title="FitPlan AI", page_icon="💪", layout="centered")
+
+# -----------------------------
+# BACKGROUND + STYLES
+# -----------------------------
+st.markdown("""
+<style>
+.stApp {
+    background: linear-gradient(to right, #ff512f, #dd2476);
+}
+
+/* Title */
+.main-title {
+    font-size: 40px;
+    font-weight: bold;
+    text-align: center;
+    color: white;
+}
+
+/* Workout Output Box */
+.workout-box {
+    background-color: #111;
+    padding: 25px;
+    border-radius: 15px;
+    border: 1px solid #444;
+    color: white;
+    font-size: 18px;
+    line-height: 1.7;
+}
+
+/* Section Title */
+.title-box {
+    font-size: 28px;
+    font-weight: bold;
+    color: #ff7b00;
+    margin-top: 20px;
+}
+</style>
+""", unsafe_allow_html=True)
+
+st.markdown('<div class="main-title">💪 FitPlan AI</div>', unsafe_allow_html=True)
+st.write("### Fitness Profile & BMI Calculator")
+
+# -----------------------------
+# LOAD MODEL (HUGGING FACE)
+# -----------------------------
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
-    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-# -------------------------
-# App Title
-# -------------------------
-
-st.title("💪 FitPlan AI - BMI Calculator & Workout Planner")
-st.write("Enter your details to get AI workout plan.")
-
-# -------------------------
-# Personal Info
-# -------------------------
-
+# -----------------------------
+# USER INPUT
+# -----------------------------
 name = st.text_input("Enter Your Name *")
 gender = st.selectbox("Gender", ["Male", "Female", "Other"])
 height_cm = st.number_input("Height (cm)", min_value=0.0)
 weight_kg = st.number_input("Weight (kg)", min_value=0.0)
 
-# -------------------------
-# Fitness Details
-# -------------------------
+st.subheader("Fitness Details")
 
 goal = st.selectbox(
     "Fitness Goal",
-    ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexible"]
+    ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexibility"]
 )
 
 equipment = st.multiselect(
     "Available Equipment",
     ["Dumbbells", "Resistance Band", "Yoga Mat", "Skipping Rope",
-     "Weight Plates", "Cycling", "Inclined Bench", "Pullups Bar", "No Equipment"]
-)
-
-fitness_level = st.radio(
-    "Fitness Level",
-    ["Beginner", "Intermediate", "Advanced"]
+     "Weight Plates", "Cycling", "Inclined Bench", "Pull-up Bar", "No Equipment"]
 )
 
-# -------------------------
-# BMI Functions
-# -------------------------
+fitness_level = st.radio("Fitness Level", ["Beginner", "Intermediate", "Advanced"])
 
+# -----------------------------
+# BMI FUNCTIONS
+# -----------------------------
 def calculate_bmi(weight, height):
     height_m = height / 100
     return round(weight / (height_m ** 2), 2)
@@ -86,11 +98,10 @@ def bmi_category(bmi):
     else:
         return "Obese"
 
-# -------------------------
-# Generate Plan
-# -------------------------
-
-if st.button("Generate Workout Plan"):
+# -----------------------------
+# SUBMIT BUTTON
+# -----------------------------
+if st.button("Generate Fitness Plan 💥"):
 
     if not name:
         st.error("Enter your name")
@@ -99,35 +110,64 @@ if st.button("Generate Workout Plan"):
     elif not equipment:
         st.error("Select equipment")
     else:
+        st.success("Profile Submitted Successfully!")
+
         bmi = calculate_bmi(weight_kg, height_cm)
         bmi_status = bmi_category(bmi)
 
-        st.success(f"BMI: {bmi} ({bmi_status})")
+        st.write(f"### 📊 BMI: {bmi} ({bmi_status})")
 
         equipment_list = ", ".join(equipment)
 
+        # -----------------------------
+        # PROMPT FOR AI
+        # -----------------------------
         prompt = f"""
-        Create a 5-day workout plan.
-
-        Name: {name}
-        Gender: {gender}
-        BMI: {bmi} ({bmi_status})
-        Goal: {goal}
-        Level: {fitness_level}
-        Equipment: {equipment_list}
-
-        Include warmup, exercises, sets, reps, rest time.
-        """
-
-        with st.spinner("Generating workout plan..."):
+Create a STRICT 5-day structured workout plan.
+
+Format EXACTLY like:
+
+Day 1:
+Warm-up:
+Exercises:
+Sets & Reps:
+Rest:
+
+Day 2:
+...
+
+User Details:
+Name: {name}
+Gender: {gender}
+BMI: {bmi} ({bmi_status})
+Goal: {goal}
+Fitness Level: {fitness_level}
+Equipment: {equipment_list}
+"""
+
+        # -----------------------------
+        # GENERATE PLAN
+        # -----------------------------
+        with st.spinner("Generating your AI workout plan..."):
             inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=400,
+                max_new_tokens=600,
                 temperature=0.7,
                 do_sample=True
             )
             result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        st.subheader("🏋️ Your Workout Plan")
-        st.write(result)
+        # -----------------------------
+        # DISPLAY OUTPUT (Styled)
+        # -----------------------------
+        st.markdown(
+            '<div class="title-box">🏋️ Your Personalized Workout Plan</div>',
+            unsafe_allow_html=True
+        )
+
+        st.markdown(f"""
+        <div class="workout-box">
+        {result.replace("Day", "<br><br><b>Day").replace("\n", "<br>")}
+        </div>
+        """, unsafe_allow_html=True)