Update app.py

app.py

@@ -8,8 +8,11 @@ model_name = "mjpsm/Positive-Affirmations-Model"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Function to extract only the affirmation
-def generate_affirmation(prompt):
+# Generation function
+def generate_affirmation(description):
+    # Structured prompt to guide model output
+    prompt = f"[SUBJECT] learning [/SUBJECT] [STREAK] current performance context [/STREAK] [CONTEXT] {description} [/CONTEXT] [AFFIRMATION]"
+
     inputs = tokenizer(prompt, return_tensors="pt")
     input_ids = inputs["input_ids"]
     
@@ -26,12 +29,16 @@ def generate_affirmation(prompt):
 
     full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract content inside [AFFIRMATION]...[/AFFIRMATION]
+    # Try to extract text between [AFFIRMATION] and [/AFFIRMATION]
     match = re.search(r"\[AFFIRMATION\](.*?)\[/AFFIRMATION\]", full_output, re.DOTALL)
     if match:
         return match.group(1).strip()
     else:
-        return "No affirmation found in the response."
+        # Fallback: try to extract everything after [AFFIRMATION]
+        fallback_match = re.search(r"\[AFFIRMATION\](.*)", full_output, re.DOTALL)
+        if fallback_match:
+            return fallback_match.group(1).strip()
+        return "⚠️ No affirmation found in the response."
 
 # Gradio interface
 demo = gr.Interface(
@@ -39,7 +46,7 @@ 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__":
@@ -48,3 +55,5 @@ if __name__ == "__main__":
 
 
 
+
+