import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load model and tokenizer from Hugging Face Hub
model_name = "mjpsm/Positive-Affirmations-Model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Generation function
def generate_affirmation(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_new_tokens=100,
            temperature=0.7,
            top_k=50,
            top_p=0.95,
            do_sample=True
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio interface
demo = gr.Interface(
    fn=generate_affirmation,
    inputs=gr.Textbox(label="Describe the player situation (e.g., 'struggled with algebra')"),
    outputs=gr.Textbox(label="AI Affirmation"),
    title="Positive Affirmation Generator",
    description="Describe a learning moment, and the model will generate a motivating affirmation."
)

if __name__ == "__main__":
    demo.launch()