Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,168 +1,83 @@
-import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
-import time
 
 # -------------------------
-# Page Config
+# 1. Function to Craft Prompt
 # -------------------------
 
-st.set_page_config(page_title="FitPlan AI", page_icon="💪")
+def create_workout_prompt(days, level, goal, equipment):
+    """
+    Creates a detailed prompt for the AI model
+    """
 
-# -------------------------
-# Load Model
-# -------------------------
+    prompt = f"""
+    Generate a {days}-day workout plan.
 
-@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
+    User Details:
+    Fitness Level: {level}
+    Goal: {goal}
+    Available Equipment: {equipment}
 
-tokenizer, model = load_model()
+    Requirements:
+    - Include warm-up
+    - Include exercises with sets and reps
+    - Include rest time
+    - Keep it structured day-wise
+    """
 
-# -------------------------
-# Title
-# -------------------------
+    return prompt
 
-st.title("💪 FitPlan AI - BMI Calculator & Workout Planner")
 
 # -------------------------
-# Personal Info
+# 2. Load Model
 # -------------------------
 
-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)
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+    return tokenizer, model
+
 
 # -------------------------
-# Fitness Details
+# 3. Generate Response
 # -------------------------
 
-goal = st.selectbox(
-    "Fitness Goal",
-    ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexible"]
-)
+def generate_workout_plan(prompt, tokenizer, model):
 
-equipment = st.multiselect(
-    "Available Equipment",
-    ["Dumbbells", "Resistance Band", "Yoga Mat", "Skipping Rope",
-     "Weight Plates", "Cycling", "Inclined Bench", "Pullups Bar", "No Equipment"]
-)
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
-fitness_level = st.radio(
-    "Fitness Level",
-    ["Beginner", "Intermediate", "Advanced"]
-)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=300,
+        temperature=0.7,
+        do_sample=True
+    )
 
-# -------------------------
-# BMI Functions
-# -------------------------
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-def calculate_bmi(weight, height):
-    height_m = height / 100
-    return round(weight / (height_m ** 2), 2)
+    return result
 
-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"
 
 # -------------------------
-# Generate Plan
+# 4. Example Usage
 # -------------------------
 
-if st.button("Generate Workout Plan"):
-
-    # Validation
-    if not name:
-        st.error("Enter your name")
-    elif height_cm <= 0 or weight_kg <= 0:
-        st.error("Enter valid height & weight")
-    elif not equipment:
-        st.error("Select equipment")
-    else:
-
-        # BMI
-        bmi = calculate_bmi(weight_kg, height_cm)
-        bmi_status = bmi_category(bmi)
-
-        st.success(f"✅ Welcome {name}! BMI: {bmi} ({bmi_status})")
-
-        # -------------------------
-        # Progress Bar Animation
-        # -------------------------
-
-        progress = st.progress(0)
-        status = st.empty()
-
-        for i in range(100):
-            time.sleep(0.01)
-            progress.progress(i + 1)
-            status.text("Generating personalized fitness plan...")
-
-        st.divider()
-
-        # -------------------------
-        # Core AI Module Layout
-        # -------------------------
-
-        st.subheader("🧠 Core AI Inference Module")
-
-        col1, col2, col3 = st.columns(3)
-
-        with col1:
-            st.info("🔽 Input Layer\nUser goals, equipment, fitness level")
-
-        with col2:
-            st.info("⚙ Processing\nAI model analysis")
-
-        with col3:
-            st.info("📤 Output Layer\nWorkout plan")
-
-        st.divider()
-
-        # -------------------------
-        # Model Input & Output
-        # -------------------------
-
-        equipment_list = ", ".join(equipment)
-
-        prompt = f"""
-        Create a structured 5-day workout plan.
-
-        User:
-        Goal: {goal}
-        Fitness Level: {fitness_level}
-        Equipment: {equipment_list}
-
-        Include exercises with sets, reps, and rest time.
-        """
+if __name__ == "__main__":
 
-        with st.spinner("AI is creating your plan..."):
-            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=400,
-                temperature=0.7,
-                do_sample=True
-            )
-            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Create prompt from user input
+    prompt = create_workout_prompt(
+        days=5,
+        level="Beginner",
+        goal="Build Muscle",
+        equipment="Dumbbells"
+    )
 
-        col4, col5 = st.columns(2)
+    print("Generated Prompt:\n", prompt)
 
-        with col4:
-            st.subheader("📥 Model Input")
-            st.write(prompt)
+    # Load model
+    tokenizer, model = load_model()
 
-        with col5:
-            st.subheader("📤 Model Output")
-            st.write(result)
+    # Generate text response
+    workout_plan = generate_workout_plan(prompt, tokenizer, model)
 
-        st.success("🏋️ Your Personalized Workout Plan Ready!")
+    print("\nAI Workout Plan:\n")
+    print(workout_plan)