Update app.py

app.py

@@ -1,14 +1,28 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
 import streamlit as st
+from huggingface_hub import InferenceClient
+import os
 
 # -------------------------
 # PAGE CONFIG
 # -------------------------
 st.set_page_config(page_title="FitPlan-AI", page_icon="💪")
-
 st.title("💪 FitPlan-AI: Personalized Fitness Profile")
 
+# -------------------------
+# LOAD HF CLIENT (CACHED)
+# -------------------------
+@st.cache_resource
+def get_hf_client():
+    hf_token = os.getenv("HF_TOKEN")
+
+    if not hf_token:
+        raise ValueError("HF_TOKEN not set in environment variables.")
+
+    return InferenceClient(
+        model="mistralai/Mistral-7B-Instruct-v0.2",
+        token=hf_token
+    )
+
 # -------------------------
 # BMI FUNCTIONS
 # -------------------------
@@ -26,6 +40,67 @@ def get_category(bmi):
     else:
         return "Obese"
 
+# -------------------------
+# PROMPT BUILDER
+# -------------------------
+def build_prompt(name, height, weight, goal, level, equipment, bmi, bmi_status):
+    equipment_list = ", ".join(equipment)
+
+    prompt = f"""
+Create a detailed 5-day personalized workout plan.
+
+Client Details:
+Name: {name}
+Height: {height} cm
+Weight: {weight} kg
+BMI: {bmi} ({bmi_status})
+Goal: {goal}
+Fitness Level: {level}
+Available Equipment: {equipment_list}
+
+Instructions:
+- Provide a structured 5-day plan.
+- Each day must include:
+  - Warm-up
+  - Main Workout (Sets x Reps)
+  - Rest time
+  - Cooldown
+- Ensure exercises match fitness level and equipment.
+- Keep it practical and realistic.
+- Format clearly as:
+
+Day 1:
+Warm-up:
+Main Workout:
+Cooldown:
+
+Continue up to Day 5.
+"""
+
+    return prompt.strip()
+
+# -------------------------
+# MODEL QUERY
+# -------------------------
+def query_model(prompt):
+    try:
+        client = get_hf_client()
+
+        response = client.chat_completion(
+            messages=[
+                {"role": "system", "content": "You are a certified professional fitness trainer."},
+                {"role": "user", "content": prompt}
+            ],
+            max_tokens=800,
+            temperature=0.7,
+            top_p=0.9
+        )
+
+        return response.choices[0].message.content.strip()
+
+    except Exception as e:
+        return f"⚠️ Model Error: {str(e)}"
+
 # -------------------------
 # FORM
 # -------------------------
@@ -78,46 +153,27 @@ if submit:
     else:
         st.success("Profile Submitted Successfully!")
 
-        # Calculate BMI
+        # BMI Calculation
         bmi = calculate_bmi(weight, height)
         bmi_status = get_category(bmi)
 
         st.write(f"### 📊 Your BMI: {bmi} ({bmi_status})")
 
-        # -------------------------
-        # GENERATE 5-DAY PLAN (SINGLE PROMPT METHOD)
-        # -------------------------
-        from prompt_builder import build_prompt
-
+        # Generate Plan
         with st.spinner("Generating your 5-day workout plan..."):
-
-            prompt, bmi, bmi_status = build_prompt(
+            prompt = build_prompt(
                 name=name,
-                gender="Not specified",  # You can add gender field later
                 height=height,
                 weight=weight,
                 goal=goal,
-                fitness_level=level,
-                equipment=equipment
-            )
-
-            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=2000,
-                temperature=0.7,
-                do_sample=True,
-                repetition_penalty=1.1
+                level=level,
+                equipment=equipment,
+                bmi=bmi,
+                bmi_status=bmi_status
             )
 
-            full_plan = tokenizer.decode(
-                outputs[0],
-                skip_special_tokens=True
-            ).strip()
+            full_plan = query_model(prompt)
 
-        # -------------------------
-        # DISPLAY PLAN
-        # -------------------------
+        # Display Plan
         st.subheader("🏋️ Your Personalized 5-Day Workout Plan")
-        st.write(full_plan)
+        st.write(full_plan)