import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel

# 🔐 Hugging Face token (stored in repo secret as HF_TOKEN)
HF_TOKEN = os.environ.get("HF_TOKEN")

# 📌 Model sources
base_model = "mistralai/Mistral-7B-v0.1"
adapter_path = "./mistral-recovery-model"  # Your uploaded LoRA adapter files

# 🚀 Load tokenizer and model in FP32 (CPU-safe)
tokenizer = AutoTokenizer.from_pretrained(base_model, use_auth_token=HF_TOKEN)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    base_model,
    torch_dtype=torch.float32,
    use_auth_token=HF_TOKEN
)

model = PeftModel.from_pretrained(model, adapter_path, use_auth_token=HF_TOKEN)

# 🧠 Pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.7)

# 🖥️ Gradio UI
def generate(prompt):
    return pipe(prompt)[0]['generated_text']

gr.Interface(
    fn=generate,
    inputs="text",
    outputs="text",
    title="🏋️ Mistral Recovery Coach (CPU Mode)",
    description="Enter your workout summary and get personalized recovery advice."
).launch()