updating the langugage model with better prompting

app.py

@@ -28,27 +28,38 @@ def run_language_model(edit_prompt, device):
     )
     tokenizer = AutoTokenizer.from_pretrained(language_model_id)
     messages = [
-        {
-            "role": "system",
-            "content": "Follow the examples and return the expected output",
-        },
+        {"role": "system", "content": "Follow the examples and return the expected output"},
         {"role": "user", "content": "swap mountain and lion"},  # example 1
         {"role": "assistant", "content": "mountain, lion"},  # example 1
         {"role": "user", "content": "change the dog with cat"},  # example 2
         {"role": "assistant", "content": "dog, cat"},  # example 2
-        {"role": "user", "content": "replace the human with a boat"},  # example 3
-        {"role": "assistant", "content": "human, boat"},  # example 3
-        {"role": "user", "content": edit_prompt},
+        {"role": "user", "content": "change the cat with a dog"},  # example 3
+        {"role": "assistant", "content": "cat, dog"},  # example 3
+        {"role": "user", "content": "replace the human with a boat"},  # example 4
+        {"role": "assistant", "content": "human, boat"},  # example 4
+        {"role": "user", "content": "in the above example change the background to the alps"},  # example 5
+        {"role": "assistant", "content": "background, alps"},  # example 5
+        {"role": "user", "content": "edit the house into a mansion"},  # example 6
+        {"role": "assistant", "content": "house, a mansion"},  # example 6
+        {"role": "user", "content": edit_prompt}
     ]
     text = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
     )
     model_inputs = tokenizer([text], return_tensors="pt").to(device)
-    generated_ids = language_model.generate(model_inputs.input_ids, max_new_tokens=512)
-    generated_ids = [
-        output_ids[len(input_ids) :]
-        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
+    with torch.no_grad():
+        generated_ids = language_model.generate(
+          model_inputs.input_ids,
+          max_new_tokens=512,
+          temperature=0.0,
+          do_sample=False
+        )
+        
+        generated_ids = [
+          output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
     response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
     to_replace, replace_with = response.split(", ")