Update app.py

app.py

@@ -1,13 +1,14 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import re
 
 # 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
+# Affirmation generator
 def generate_affirmation(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
     input_ids = inputs["input_ids"]
@@ -23,17 +24,14 @@ def generate_affirmation(prompt):
             pad_token_id=tokenizer.eos_token_id
         )
 
-    # Decode full output
     full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Remove the input prompt from the output to isolate the generated part
-    if full_output.startswith(prompt):
-        affirmation = full_output[len(prompt):].strip()
+    # Extract [AFFIRMATION] ... [/AFFIRMATION] only
+    match = re.search(r"\[AFFIRMATION\](.*?)\[/AFFIRMATION\]", full_output, re.DOTALL)
+    if match:
+        return match.group(1).strip()
     else:
-        affirmation = full_output.strip()
-
-    return affirmation
-
+        return "⚠️ No [AFFIRMATION] section found in the response."
 
 # Gradio interface
 demo = gr.Interface(
@@ -41,9 +39,11 @@ demo = gr.Interface(
     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."
+    description="Describe a learning moment, and receive an uplifting affirmation generated by AI."
 )
 
 if __name__ == "__main__":
     demo.launch()
 
+
+