hit the token limit

app.py

@@ -8,6 +8,7 @@ import torchaudio.transforms as T
 import re
 import logging, sys
 import json
+from llama_cpp import Llama
 
 
 logging.basicConfig(stream=sys.stdout, level=logging.INFO, force=True)
@@ -51,6 +52,25 @@ SYSTEM_PROMPT = (
     "If you need to ask the user a follow-up question, do so clearly.\n",
 )
 
+def truncate_prompt(prompt, max_tokens=2048):
+    # Use your model's tokenizer here; this is a placeholder
+    tokens = prompt.split()  # Replace with actual tokenization
+    if len(tokens) > max_tokens:
+        tokens = tokens[:max_tokens]
+    return " ".join(tokens)
+
+# Initialize your model (adjust path and params as needed)
+llm = Llama(model_path=model_path)
+
+def truncate_prompt_llama(prompt, max_tokens=2048):
+    # Tokenize the prompt using llama_cpp's tokenizer
+    tokens = llm.tokenize(prompt.encode("utf-8"))
+    if len(tokens) > max_tokens:
+        # Truncate tokens and decode back to string
+        tokens = tokens[:max_tokens]
+        prompt = llm.detokenize(tokens).decode("utf-8", errors="ignore")
+    return prompt
+
 # ========== Generator handler ==========
 def on_submit(symptoms_text, history):
     log = []
@@ -103,8 +123,8 @@ def on_submit(symptoms_text, history):
         f"{SYSTEM_PROMPT}",
         f"User symptoms: '{cleaned}'",
         f"Relevant ICD-10 context:\n{context_text}",
-        "Respond with your top 3 ICD-10 codes and their confidence scores.",
     ])
+    prompt = truncate_prompt_llama(prompt, max_tokens=2048)
 
     msg = "✏️ Prompt built"
     log.append(msg)