speedup

app.py

@@ -1,10 +1,15 @@
 import gradio as gr
 import spaces
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load the model and tokenizer (you may need to adjust device_map or other settings depending on your hardware)
+# Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("TheBloke/Chronoboros-33B-GPTQ")
 model = AutoModelForCausalLM.from_pretrained("TheBloke/Chronoboros-33B-GPTQ", device_map="auto")
+model.eval()  # set model to evaluation mode
+
+# Optional: Use torch.compile() if you're on PyTorch 2.0+ for further speed-up
+# model = torch.compile(model)
 
 @spaces.GPU
 def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
@@ -17,21 +22,26 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
             prompt += f"Assistant: {assistant_text}\n"
     prompt += f"User: {message}\nAssistant: "
 
-    # Tokenize the prompt and generate a response
+    # Tokenize the prompt
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
-    output_ids = model.generate(
-        input_ids,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True,  # enable sampling for varied responses
-    )
-    # Get only the newly generated tokens (after the prompt)
+    
+    # Generate the response with no gradients
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+        )
+    
+    # Extract the new tokens
     new_tokens = output_ids[0][input_ids.shape[1]:]
-
-    # Simulate streaming by yielding partial responses token by token
-    for i in range(new_tokens.shape[0]):
-        current_response = tokenizer.decode(new_tokens[: i + 1], skip_special_tokens=True)
+    
+    # Stream output in chunks (e.g., 5 tokens per chunk)
+    chunk_size = 5
+    for i in range(0, new_tokens.shape[0], chunk_size):
+        current_response = tokenizer.decode(new_tokens[: i + chunk_size], skip_special_tokens=True)
         yield current_response
 
 # Configure the ChatInterface with additional inputs