Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,9 +5,10 @@ import torch
 st.set_page_config(page_title="FitPlan AI", layout="centered")
 
 # LOAD MODEL 
+@st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
+    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
     return tokenizer, model
 
 tokenizer, model = load_model()
@@ -147,31 +148,41 @@ if st.button(" Submit Profile"):
         equipment_list = ", ".join(equipment)
 
         prompt = f"""
-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}
+Instruction:
+Generate a structured 5-day workout plan.
+
+User:
+Gender: {gender}
+BMI: {bmi:.2f} ({bmi_status})
+Goal: {goal}
+Level: {fitness_level}
+Equipment: {equipment_list}
+
+Rules:
+- Only 5 days.
+- Each day must include:
+  Focus
+  4 Exercises
+  Sets x Reps
+  Rest
+- No repetition.
+- No extra explanation.
 
 Start directly with:
-
 Day 1:
 """
 
+
         with st.spinner("Generating your AI workout plan..."):
 
             inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=900,
-                temperature=0.7,
-                do_sample=True
+                max_new_tokens=400,
+                num_beams=4,
+                early_stopping=True,
+                no_repeat_ngram_size=3
             )
 
             result = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()