Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,34 +1,34 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-# Load model only once
-@st.cache_resource
-def load_model():
-    return pipeline("text-generation", model="google/flan-t5-base")
+# -------------------------
+# Page Config
+# -------------------------
 
-generator = load_model()
+st.set_page_config(page_title="FitPlan AI", page_icon="💪")
 
-st.set_page_config(page_title="FitPlan AI - BMI Calculator", page_icon="💪")
-st.title("💪 FitPlan AI - Fitness Profile & BMI Calculator")
+st.title("💪 FitPlan AI — Personalized Workout Generator")
 
-st.write("Fill in your details to calculate your BMI and get AI workout plan.")
+st.write("Fill your details to generate AI workout plan.")
 
 # -------------------------
-# 1. Personal Information
+# Load Model (Small for HF Spaces)
 # -------------------------
 
-name = st.text_input("Enter Your Name *")
-
-gender = st.selectbox("Select Gender", ["Male", "Female", "Other"])
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+    return tokenizer, model
 
-height_cm = st.number_input("Enter Height (in centimeters) *", min_value=0.0, format="%.2f")
-weight_kg = st.number_input("Enter Weight (in kilograms) *", min_value=0.0, format="%.2f")
+tokenizer, model = load_model()
 
 # -------------------------
-# 2. Fitness Details
+# User Inputs
 # -------------------------
 
-st.subheader("Fitness Details")
+name = st.text_input("Enter Your Name")
 
 goal = st.selectbox(
     "Select Your Fitness Goal",
@@ -46,56 +46,76 @@ fitness_level = st.radio(
 )
 
 # -------------------------
-# BMI Functions
+# Function: Create Prompt
 # -------------------------
 
-def calculate_bmi(weight, height_cm):
-    height_m = height_cm / 100
-    bmi = weight / (height_m ** 2)
-    return round(bmi, 2)
-
-def bmi_category(bmi):
-    if bmi < 18.5:
-        return "Underweight"
-    elif bmi < 24.9:
-        return "Normal"
-    elif bmi < 29.9:
-        return "Overweight"
-    else:
-        return "Obese"
+def create_prompt(goal, equipment, fitness_level):
+
+    equipment_list = ", ".join(equipment)
+
+    # Goal instructions
+    goal_map = {
+        "Build Muscle": "Focus on muscle hypertrophy training.",
+        "Weight Loss": "Focus on fat burning and cardio workouts.",
+        "Strength Gain": "Focus on heavy compound strength exercises.",
+        "Abs Building": "Focus on core and abdominal exercises.",
+        "Flexible": "Focus on stretching and mobility exercises."
+    }
+
+    # Level instructions
+    level_map = {
+        "Beginner": "Use simple low intensity exercises.",
+        "Intermediate": "Use moderate intensity exercises.",
+        "Advanced": "Use high intensity advanced exercises."
+    }
+
+    prompt = f"""
+    Generate a 5-day workout plan.
+
+    Goal: {goal}
+    Fitness Level: {fitness_level}
+    Equipment: {equipment_list}
+
+    Goal Instruction:
+    {goal_map[goal]}
+
+    Level Instruction:
+    {level_map[fitness_level]}
+
+    Requirements:
+    - Include warm-up
+    - Exercises with sets and reps
+    - Rest time
+    - Day-wise structure
+    """
+
+    return prompt
 
 # -------------------------
-# Submit Button
+# Generate Plan
 # -------------------------
 
-if st.button("Submit Profile"):
+if st.button("Generate Workout Plan"):
 
-    # Validation
     if not name:
-        st.error("Please enter your name.")
-    elif height_cm <= 0 or weight_kg <= 0:
-        st.error("Enter valid height and weight.")
+        st.error("Enter your name")
     elif not equipment:
-        st.error("Select at least one equipment option.")
+        st.error("Select at least one equipment")
     else:
-        st.success("✅ Profile Submitted Successfully!")
-
-        # Calculate BMI
-        bmi = calculate_bmi(weight_kg, height_cm)
-        bmi_status = bmi_category(bmi)
-
-        st.write(f"### 📊 BMI: {bmi} ({bmi_status})")
+        prompt = create_prompt(goal, equipment, fitness_level)
 
-        equipment_list = ", ".join(equipment)
+        with st.spinner("Generating your AI workout plan..."):
 
-        prompt = f"""
-        Generate a 5-day structured workout plan.
+            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
-        Include warmup, exercises with sets & reps, rest time.
-        """
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=300,
+                temperature=0.7,
+                do_sample=True
+            )
 
-        with st.spinner("Generating your AI workout plan..."):
-            result = generator(prompt, max_new_tokens=400)[0]["generated_text"]
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        st.subheader("🏋️ Your Personalized Workout Plan")
+        st.success(f"Workout Plan for {name}")
         st.write(result)