Upload interactive_inference.py

interactive_inference.py

@@ -65,7 +65,10 @@ def interactive_session():
             # We enable modulation to see the effect of the trained controller
             # The controller predicts modulation based on the input prompt
             
-            inputs = model.tokenizer(user_input, return_tensors="pt").to(device)
+            # Format the prompt to match training distribution
+            prompt = f"User: {user_input}\nModel: "
+            
+            inputs = model.tokenizer(prompt, return_tensors="pt").to(device)
             
             with torch.no_grad():
                 # 1. Predict Modulation
@@ -78,17 +81,22 @@ def interactive_session():
                     **inputs,
                     max_new_tokens=128,
                     do_sample=True,
-                    temperature=0.7,
-                    repetition_penalty=1.1,
+                    temperature=0.6,
+                    top_p=0.9,
+                    repetition_penalty=1.2,
                     pad_token_id=model.tokenizer.eos_token_id
                 )
                 
                 model.clear_modulation()
                 
             response = model.tokenizer.decode(out_ids[0], skip_special_tokens=True)
-            # Strip prompt if included (Gemma usually includes it)
-            if response.startswith(user_input):
-                response = response[len(user_input):].strip()
+            
+            # Clean up response (Remove the prompt part)
+            if response.startswith(prompt):
+                response = response[len(prompt):].strip()
+            elif "Model:" in response:
+                response = response.split("Model:")[-1].strip()
+
                 
             print(f"MODEL: {response}")
             print(f"   [Modulation Norm: {torch.norm(modulation).item():.2f}]")