Update app.py

app.py

@@ -10,16 +10,30 @@ model = AutoModelForCausalLM.from_pretrained(model_name)
 # Generation function
 def generate_affirmation(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"]
+    
     with torch.no_grad():
         outputs = model.generate(
-            inputs["input_ids"],
+            input_ids,
             max_new_tokens=100,
             temperature=0.7,
             top_k=50,
             top_p=0.95,
-            do_sample=True
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # 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()
+    else:
+        affirmation = full_output.strip()
+
+    return affirmation
+
 
 # Gradio interface
 demo = gr.Interface(