Better performance and config?

app.py

@@ -5,14 +5,15 @@ MODEL_ID = "TinyLlama/TinyLlama_v1.1_math_code"
 
 handler = loop.EndpointHandler(MODEL_ID)
 
-def respond(prompt, grammar):
-    args = { "inputs": prompt, "grammar": grammar }
+def respond(prompt, grammar, max_new_tokens, max_time):
+    args = { "inputs": prompt, "grammar": grammar, "max-new-tokens": max_new_tokens, "max-time": max_time }
     return handler(args)[0]
 
 demo = gr.Interface(
     respond,
     inputs=["textarea", "textarea"],
-    outputs=["textarea"]
+    outputs=["textarea"],
+    additional_inputs=[gr.Number(value=512, precision=0), gr.Number(value=30, precision=0)]
 )
 
 if __name__ == "__main__":

loop.py

@@ -1,23 +1,51 @@
 from typing import Dict, List, Any
 
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from transformers.generation.logits_process import LogitsProcessorList, InfNanRemoveLogitsProcessor
 from transformers_gad.grammar_utils import IncrementalGrammarConstraint
 from transformers_gad.generation.logits_process import GrammarAlignedOracleLogitsProcessor
 
+def safe_int_cast(str, default):
+    try:
+        return int(str)
+    except ValueError:
+        return default
+
 class EndpointHandler():
     def __init__(self, path=""):
         # Preload
+        DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+        DTYPE = torch.bfloat16
+
+        self.device = torch.device(DEVICE)
+
         self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
         self.model = AutoModelForCausalLM.from_pretrained(path)
+        self.model.to(self.device)
+        self.model.to(dtype=DTYPE)
+        self.model.resize_token_embeddings(len(self.tokenizer))
+        self.model = torch.compile(self.model, mode='reduce-overhead', fullgraph=True)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         # do it!
-        inputs = data.get("inputs",data)
+        MAX_NEW_TOKENS=512
+        MAX_TIME=30
+        TEMPERATURE = 1.0
+        REPETITION_PENALTY = 1.0
+        TOP_P = 1.0
+        TOP_K = 0
+
+        inputs = data.get("inputs", data)
         grammar_str = data.get("grammar", "")
-        MAX_NEW_TOKENS=4096
-        MAX_TIME=300
+        max_new_tokens = safe_int_cast(data.get("max-new-tokens"), MAX_NEW_TOKENS)
+        max_time = safe_int_cast(data.get("max-time"), MAX_TIME)
+
+        print("=== GOT GRAMMAR ===")
         print(grammar_str)
+        print("===================")
         grammar = IncrementalGrammarConstraint(grammar_str, "root", self.tokenizer)
 
         # Initialize logits processor for the grammar
@@ -28,14 +56,24 @@ class EndpointHandler():
             gad_oracle_processor,
         ])
 
-        input_ids = self.tokenizer([inputs], add_special_tokens=False, return_tensors="pt")["input_ids"]
+        input_ids = self.tokenizer([inputs], add_special_tokens=False, return_tensors="pt", padding=True)["input_ids"]
+        input_ids = input_ids.to(self.model.device)
 
         output = self.model.generate(
                     input_ids,
                     do_sample=True,
-                    max_time=MAX_TIME,
-                    max_new_tokens=MAX_NEW_TOKENS,
-                    logits_processor=logits_processors
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    max_time=max_time,
+                    max_new_tokens=max_new_tokens,
+                    top_p=TOP_P,
+                    top_k=TOP_K,
+                    repetition_penalty=REPETITION_PENALTY,
+                    temperature=TEMPERATURE,
+                    logits_processor=logits_processors,
+                    num_return_sequences=1,
+                    return_dict_in_generate=True,
+                    output_scores=True
                 )
         
         gad_oracle_processor.reset()