src/streamlit_app.py

import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# -------------------------
# 1. Function to Craft Prompt
# -------------------------

def create_workout_prompt(days, level, goal, equipment):
    """
    Creates a detailed prompt for the AI model
    """

    prompt = f"""
    Generate a {days}-day workout plan.

    User Details:
    Fitness Level: {level}
    Goal: {goal}
    Available Equipment: {equipment}

    Requirements:
    - Include warm-up
    - Include exercises with sets and reps
    - Include rest time
    - Keep it structured day-wise
    """

    return prompt


# -------------------------
# 2. Load Model
# -------------------------

def load_model():
    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
    return tokenizer, model


# -------------------------
# 3. Generate Response
# -------------------------

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


# -------------------------
# 4. Example Usage
# -------------------------

if __name__ == "__main__":

    # Create prompt from user input
    prompt = create_workout_prompt(
        days=5,
        level="Beginner",
        goal="Build Muscle",
        equipment="Dumbbells"
    )

    print("Generated Prompt:\n", prompt)

    # Load model
    tokenizer, model = load_model()

    # Generate text response
    workout_plan = generate_workout_plan(prompt, tokenizer, model)

    print("\nAI Workout Plan:\n")
    print(workout_plan)