Upload gamemaster.py

game/gamemaster.py

@@ -189,15 +189,19 @@ def _run_model(system: str, user: str) -> str:
 
     messages = [{"role": "system", "content": system},
                 {"role": "user", "content": user}]
-    inputs = _tokenizer.apply_chat_template(
-        messages, add_generation_prompt=True, return_tensors="pt"
+    # return_dict=True gives a BatchEncoding (input_ids + attention_mask) on
+    # every transformers version — newer releases return it by default, and
+    # passing it positionally to generate() crashes on `.shape`.
+    enc = _tokenizer.apply_chat_template(
+        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
     ).to(_model.device)
     with torch.no_grad():
         out = _model.generate(
-            inputs, max_new_tokens=400, do_sample=True, temperature=0.7, top_p=0.9,
+            **enc, max_new_tokens=400, do_sample=True, temperature=0.7, top_p=0.9,
             pad_token_id=_tokenizer.eos_token_id,
         )
-    text = _tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)
+    n_in = enc["input_ids"].shape[1]
+    text = _tokenizer.decode(out[0][n_in:], skip_special_tokens=True)
     return text