Update chatcode/runtime.py

chatcode/runtime.py

@@ -61,8 +61,6 @@ def load_engine(adapter: str | None = None):
 
 def answer(tokenizer, model, prompt: str, history: list[dict] | None = None,
            max_new_tokens: int = CONFIG["max_new_tokens"]) -> str:
-    import torch
-
     if not 1 <= max_new_tokens <= 2048:
         raise ValueError("Response length must be between 1 and 2,048 tokens.")
     messages = conversation(prompt, history)
@@ -75,11 +73,24 @@ def answer(tokenizer, model, prompt: str, history: list[dict] | None = None,
         if len(messages) <= 2:
             raise ValueError("This message is too long. Please send a smaller code example.")
         del messages[1:3]
+    return generate_messages(tokenizer, model, messages, max_new_tokens=max_new_tokens)
+
+
+def generate_messages(tokenizer, model, messages, max_new_tokens=768,
+                      context_limit=4096, do_sample=True):
+    import torch
+    if not 1 <= max_new_tokens <= 2048 or not 256 <= context_limit <= 8192:
+        raise ValueError("Invalid response or context length.")
+    inputs = tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+        return_dict=True, return_tensors="pt")
+    if inputs["input_ids"].shape[1] + max_new_tokens > context_limit:
+        raise ValueError("Source packet is too long. Shorten the excerpts; no sources were silently removed.")
     inputs = inputs.to(model.device)
+    sampling = {"temperature": 0.7, "top_p": 0.8, "top_k": 20} if do_sample else {}
     with torch.inference_mode():
         output = model.generate(
-            **inputs, max_new_tokens=max_new_tokens, do_sample=True,
-            temperature=0.7, top_p=0.8, top_k=20,
+            **inputs, max_new_tokens=max_new_tokens, do_sample=do_sample, **sampling,
             pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)
     reply = tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
     if not reply: