Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,23 +1,32 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
+import requests
 
 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-small")
-    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
-    return tokenizer, model
+# ---------------------------------------------------
+# HUGGING FACE API CONFIG
+# ---------------------------------------------------
 
-tokenizer, model = load_model()
+API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
+headers = {
+    "Authorization": f"Bearer {st.secrets['HF_TOKEN']}"
+}
 
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()
+
+# ---------------------------------------------------
 # TITLE
-st.title(" FitPlan AI – User Fitness Profile")
+# ---------------------------------------------------
+
+st.title("🏋️ FitPlan AI – User Fitness Profile")
 
+# ---------------------------------------------------
 # PERSONAL INFORMATION
-st.subheader(" Personal Information")
+# ---------------------------------------------------
+
+st.subheader("👤 Personal Information")
 
 name = st.text_input("Enter Your Name")
 
@@ -30,6 +39,7 @@ gender = st.radio(
 # ---------------------------------------------------
 # HEIGHT & WEIGHT
 # ---------------------------------------------------
+
 col1, col2 = st.columns(2)
 
 with col1:
@@ -50,7 +60,10 @@ with col2:
         step=0.1
     )
 
-#BMI Function
+# ---------------------------------------------------
+# BMI FUNCTION
+# ---------------------------------------------------
+
 def bmi_category(bmi):
     if bmi < 18.5:
         return "Underweight"
@@ -61,7 +74,6 @@ def bmi_category(bmi):
     else:
         return "Obese"
 
-# BMI CALCULATION
 bmi = None
 
 if height > 0 and weight > 0:
@@ -71,7 +83,10 @@ if height > 0 and weight > 0:
     st.metric("📊 Your BMI", f"{bmi:.2f}")
     st.info(f"BMI Category: {bmi_category(bmi)}")
 
+# ---------------------------------------------------
 # FITNESS GOAL
+# ---------------------------------------------------
+
 st.subheader("🎯 Fitness Goal")
 
 goal = st.selectbox(
@@ -85,9 +100,11 @@ goal = st.selectbox(
     ]
 )
 
+# ---------------------------------------------------
 # EQUIPMENT
+# ---------------------------------------------------
 
-st.subheader(" Available Equipment")
+st.subheader("🏋️ Available Equipment")
 
 equipment_map = {}
 
@@ -119,8 +136,10 @@ with col3:
 
 equipment = [item for item, selected in equipment_map.items() if selected]
 
-
+# ---------------------------------------------------
 # FITNESS LEVEL
+# ---------------------------------------------------
+
 st.subheader("📈 Fitness Level")
 
 fitness_level = st.radio(
@@ -129,8 +148,11 @@ fitness_level = st.radio(
     horizontal=True
 )
 
+# ---------------------------------------------------
 # SUBMIT BUTTON
-if st.button(" Submit Profile"):
+# ---------------------------------------------------
+
+if st.button("🚀 Submit Profile"):
 
     if not name:
         st.error("Please enter your name.")
@@ -147,49 +169,36 @@ if st.button(" Submit Profile"):
         bmi_status = bmi_category(bmi)
         equipment_list = ", ".join(equipment)
 
-        prompt = f"""
+        prompt = f"""<s>[INST]
 You are a certified professional fitness trainer.
 
 Generate a structured 5-day workout plan.
 
-You MUST follow the format EXACTLY.
-Do NOT add any extra text.
-Do NOT add explanations.
-Do NOT change spacing.
-Do NOT remove indentation.
-
 Return ONLY the workout plan.
-
-Strict Format Example:
-
-Day 1: Upper Body
-1. Push-Ups
-   Sets: 3
-   Reps: 10-12
-   Rest: 60
-
-Now generate for:
+Do not add explanations.
 
 Gender: {gender}
 BMI: {bmi:.2f} ({bmi_status})
 Goal: {goal}
 Fitness Level: {fitness_level}
 Equipment: {equipment_list}
+[/INST]
 """
 
         with st.spinner("Generating your AI workout plan..."):
 
-            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-            
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=600,
-                temperature=0.3,
-                do_sample=False,
-                repetition_penalty=1.2
-            )
-            
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
+            output = query({
+                "inputs": prompt,
+                "parameters": {
+                    "max_new_tokens": 600,
+                    "temperature": 0.3
+                }
+            })
+
+            if isinstance(output, list):
+                response = output[0]["generated_text"]
+            else:
+                response = "Error generating plan. Please try again."
+
         st.subheader("🏋️ Your Personalized Workout Plan")
-        st.write(response)
+        st.markdown(response)