Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,6 @@
 import streamlit as st
-import transformers 
-from transformers import pipeline 
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM  
+import torch 
 
 # --- Page Config ---
 st.set_page_config(page_title="FitPlan AI", page_icon="🏋", layout="wide")
@@ -23,26 +23,26 @@ if 'user_data' not in st.session_state:
 if 'profile_complete' not in st.session_state:
    st.session_state.profile_complete = False
 
-# ---  Functions ---
+# --- Helper Functions ---
 def calculate_bmi(w, h_cm):
    if h_cm > 0:
        return round(w / ((h_cm / 100) ** 2), 2)
-   return None
+   return 0
 
 def bmi_category(bmi):
-   if bmi is None: return "Unknown"
    if bmi < 18.5: return "Underweight"
    if bmi < 25: return "Normal"
    if bmi < 30: return "Overweight"
    return "Obese"
 
-# - ---
+
 @st.cache_resource 
 def load_model():
-    # Using text-generation as per the environment's supported list
-    return pipeline("text-generation", model="google/flan-t5-base")
+    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+    return tokenizer, model
 
-generator = load_model() 
+tokenizer, model = load_model()
 
 # --- Sidebar Navigation ---
 with st.sidebar:
@@ -57,29 +57,33 @@ if menu == "👤 Profile":
    name = st.text_input("Name", value=st.session_state.user_data['name'])
    gender = st.selectbox("Gender", ["Male", "Female", "Other"])
    col1, col2 = st.columns(2)
-   u_height = col1.number_input("Height (cm)", min_value=0.0, value=st.session_state.user_data['height'])
-   u_weight = col2.number_input("Weight (kg)", min_value=0.0, value=st.session_state.user_data['weight'])
+   height = col1.number_input("Height (cm)", min_value=0.0, value=st.session_state.user_data['height'])
+   weight = col2.number_input("Weight (kg)", min_value=0.0, value=st.session_state.user_data['weight'])
    
    goal = st.selectbox("Fitness Goal", ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexible"])
    equipment = st.multiselect("Equipment", ["Dumbbells", "Barbell", "Kettlebell", "Resistance Band", "Yoga Mat", "Full Gym", "No Equipment"])
    fitness_level = st.select_slider("Fitness Level", options=["Beginner", "Intermediate", "Advanced"])
 
-   bmi = calculate_bmi(u_weight, u_height)
+   bmi = calculate_bmi(weight, height)
 
-   #  ---
-   if st.button("Submit Profile"):
+  
+   if st.button(" Submit Profile"):
     
     if not name:
         st.error("Please enter your name.")
+ 
+    elif height <= 0 or weight <= 0:
+        st.error("Please enter valid height and weight.")
+ 
     elif not equipment:
         st.error("Please select at least one equipment option.")
-    elif bmi is None:
-        st.error("Please enter valid height and weight.")
+ 
     else:
-        st.success("✅ Profile Submitted Successfully!")
+        st.success(" Profile Submitted Successfully!")
         
+        # Syncing data for dashboard
         st.session_state.user_data.update({
-            'name': name, 'height': u_height, 'weight': u_weight, 
+            'name': name, 'height': height, 'weight': weight, 
             'goal': goal, 'equip': equipment, 'level': fitness_level, 'gender': gender
         })
         st.session_state.profile_complete = True
@@ -87,31 +91,37 @@ if menu == "👤 Profile":
         bmi_status = bmi_category(bmi)
         equipment_list = ", ".join(equipment)
  
-        # Prompt
         prompt = f"""
-        Generate a 5-day structured workout plan.
- 
-        User Details:
-        Name: {name}
-        Gender: {gender}
-        BMI: {bmi:.2f} ({bmi_status})
-        Goal: {goal}
-        Fitness Level: {fitness_level}
-        Available Equipment: {equipment_list}
- 
-        Requirements:
-        - Include warmup
-        - Include exercises with sets and reps
-        - Include rest time
-        - Adjust intensity based on BMI and fitness level
-        - Keep it structured day-wise
-        """
+You are a certified professional fitness trainer.
+
+Create a detailed 5-day workout plan.
+
+User Information:
+- Gender: {gender}
+- BMI: {bmi:.2f} ({bmi_status})
+- Goal: {goal}
+- Fitness Level: {fitness_level}
+- Equipment Available: {equipment_list}
+
+Start directly with:
+
+Day 1:
+"""
  
         with st.spinner("Generating your AI workout plan..."):
             
-            result = generator(prompt, max_new_tokens=400)[0]["generated_text"]
+            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+ 
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=900,
+                temperature=0.7,
+                do_sample=True
+            )
+ 
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
  
-        st.subheader("🏋️ Your Personalized Workout Plan")
+        st.subheader(" Your Personalized Workout Plan")
         st.write(result)
 
 elif menu == "📊 Dashboard":
@@ -123,4 +133,4 @@ elif menu == "📊 Dashboard":
        st.title(f"Welcome, {ud['name']}!")
        st.metric("Current BMI", f"{bmi}")
        st.write(f"**Current Goal:** {ud['goal']}")
-       st.write(f"**Equipment:** {', '.join(ud['equip'])}")
+       st.write(f"**Equipment Available:** {', '.join(ud['equip'])}")