Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,9 +6,6 @@ import torch
 # -------------------------
 
 def create_workout_prompt(days, level, goal, equipment):
-    """
-    Creates a detailed prompt for the AI model
-    """
 
     prompt = f"""
     Generate a {days}-day workout plan.
@@ -29,56 +26,56 @@ def create_workout_prompt(days, level, goal, equipment):
 
 
 # -------------------------
-# 2. Load Model
+# 2. Load Model (Cached)
 # -------------------------
 
+@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
 
+tokenizer, model = load_model()
+
 
 # -------------------------
-# 3. Generate Response
+# Streamlit UI
 # -------------------------
 
-def generate_workout_plan(prompt, tokenizer, model):
-
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=300,
-        temperature=0.7,
-        do_sample=True
-    )
-
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return result
+st.title("💪 FitPlan AI — Workout Generator")
 
+days = st.selectbox("Number of Days", [3, 4, 5, 6])
+level = st.selectbox("Fitness Level", ["Beginner", "Intermediate", "Advanced"])
+goal = st.selectbox("Goal", ["Build Muscle", "Weight Loss", "Strength Gain"])
+equipment = st.text_input("Equipment (e.g., Dumbbells)")
 
 # -------------------------
-# 4. Example Usage
+# Generate Button
 # -------------------------
 
-if __name__ == "__main__":
+if st.button("Generate Workout Plan"):
+
+    if not equipment:
+        st.error("Please enter equipment")
+    else:
+        prompt = create_workout_prompt(days, level, goal, equipment)
+
+        st.write("### Generated Prompt")
+        st.code(prompt)
 
-    # Create prompt from user input
-    prompt = create_workout_prompt(
-        days=5,
-        level="Beginner",
-        goal="Build Muscle",
-        equipment="Dumbbells"
-    )
+        with st.spinner("Generating plan..."):
+            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
-    print("Generated Prompt:\n", prompt)
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=300,
+                temperature=0.7,
+                do_sample=True
+            )
 
-    # Load model
-    tokenizer, model = load_model()
+            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Generate text response
-    workout_plan = generate_workout_plan(prompt, tokenizer, model)
+        st.success("Workout Plan Generated ✅")
 
-    print("\nAI Workout Plan:\n")
-    print(workout_plan)
+        st.subheader("🏋️ Your Workout Plan")
+        st.write(result)