import gradio as gr
from transformers import LlamaForCausalLM, LlamaTokenizer
import torch

# Load tokenizer
tokenizer = LlamaTokenizer.from_pretrained("prashb27/Llama-2-7b-chat-finetune-gym1")

# Load the model
model = LlamaForCausalLM.from_pretrained(
    "prashb27/Llama-2-7b-chat-finetune-gym1", 
    device_map="cpu",  # Force it to run on CPU
)

# Define the inference function
def generate_workout_plan(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to("cpu")
    outputs = model.generate(**inputs, max_new_tokens=50)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Define Gradio interface
def gradio_interface(user_input):
    return generate_workout_plan(user_input)

# Gradio UI layout
iface = gr.Interface(
    fn=gradio_interface,  # Function to generate workout plan
    inputs=gr.Textbox(lines=2, placeholder="Enter your query..."),  # User input
    outputs="text",  # Output is text
    title="Workout Plan Generator",  # Title for the app
    description="Enter your workout query to generate a personalized plan.",  # Description of the app
)

# Launch the app
iface.launch()