Trying to work out concurency

app.py

@@ -1,4 +1,3 @@
-'''
 import os
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -6,6 +5,7 @@ import torch
 
 model_name_2_7B_instruct = "Zyphra/Zamba2-2.7B-instruct"
 model_name_7B_instruct = "Zyphra/Zamba2-7B-instruct"
+max_context_length = 4096
 
 tokenizer_2_7B_instruct = AutoTokenizer.from_pretrained(model_name_2_7B_instruct)
 model_2_7B_instruct = AutoModelForCausalLM.from_pretrained(
@@ -41,7 +41,16 @@ def generate_response_2_7B_instruct(chat_history, max_new_tokens):
             sample.append({'role': 'assistant', 'content': turn[1]})
     chat_sample = tokenizer_2_7B_instruct.apply_chat_template(sample, tokenize=False)
     input_ids = tokenizer_2_7B_instruct(chat_sample, return_tensors='pt', add_special_tokens=False).to(model_2_7B_instruct.device)
-    outputs = model_2_7B_instruct.generate(**input_ids, max_new_tokens=int(max_new_tokens), return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
+
+    max_new_tokens = int(max_new_tokens)
+    max_input_length = max_context_length - max_new_tokens
+    if input_ids['input_ids'].size(1) > max_input_length:
+        input_ids['input_ids'] = input_ids['input_ids'][:, -max_input_length:]
+        if 'attention_mask' in input_ids:
+            input_ids['attention_mask'] = input_ids['attention_mask'][:, -max_input_length:]
+
+    with torch.no_grad():
+        outputs = model_2_7B_instruct.generate(**input_ids, max_new_tokens=int(max_new_tokens), return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
     """
     outputs = model_2_7B_instruct.generate(
         input_ids=input_ids,
@@ -59,6 +68,11 @@ def generate_response_2_7B_instruct(chat_history, max_new_tokens):
     """
     generated_text = tokenizer_2_7B_instruct.decode(outputs[0])
     assistant_response = extract_assistant_response(generated_text)
+
+    del input_ids
+    del outputs
+    torch.cuda.empty_cache()
+
     return assistant_response
 
 def generate_response_7B_instruct(chat_history, max_new_tokens):
@@ -70,7 +84,16 @@ def generate_response_7B_instruct(chat_history, max_new_tokens):
             sample.append({'role': 'assistant', 'content': turn[1]})
     chat_sample = tokenizer_7B_instruct.apply_chat_template(sample, tokenize=False)
     input_ids = tokenizer_7B_instruct(chat_sample, return_tensors='pt', add_special_tokens=False).to(model_7B_instruct.device)
-    outputs = model_7B_instruct.generate(**input_ids, max_new_tokens=int(max_new_tokens), return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
+
+    max_new_tokens = int(max_new_tokens)
+    max_input_length = max_context_length - max_new_tokens
+    if input_ids['input_ids'].size(1) > max_input_length:
+        input_ids['input_ids'] = input_ids['input_ids'][:, -max_input_length:]
+        if 'attention_mask' in input_ids:
+            input_ids['attention_mask'] = input_ids['attention_mask'][:, -max_input_length:]
+
+    with torch.no_grad():
+        outputs = model_7B_instruct.generate(**input_ids, max_new_tokens=int(max_new_tokens), return_dict_in_generate=False, output_scores=False, use_cache=True, num_beams=1, do_sample=False)
     """
     outputs = model_7B_instruct.generate(
         input_ids=input_ids,
@@ -88,6 +111,11 @@ def generate_response_7B_instruct(chat_history, max_new_tokens):
     """
     generated_text = tokenizer_7B_instruct.decode(outputs[0])
     assistant_response = extract_assistant_response(generated_text)
+
+    del input_ids
+    del outputs
+    torch.cuda.empty_cache()
+
     return assistant_response
 
 with gr.Blocks() as demo:
@@ -222,18 +250,23 @@ def generate_response_2_7B_instruct(chat_history, max_new_tokens):
     if input_ids.size(1) > max_input_length:
         input_ids = input_ids[:, -max_input_length:]  # Truncate from the left (oldest tokens)
 
-    outputs = model_2_7B_instruct.generate(
-        input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
-        return_dict_in_generate=False,
-        output_scores=False,
-        use_cache=True,
-        num_beams=1,
-        do_sample=False
-    )
+    with torch.no_grad():
+        outputs = model_2_7B_instruct.generate(
+            input_ids=input_ids,
+            max_new_tokens=max_new_tokens,
+            return_dict_in_generate=False,
+            output_scores=False,
+            use_cache=True,
+            num_beams=1,
+            do_sample=False
+        )
 
     generated_text = tokenizer_2_7B_instruct.decode(outputs[0])
     assistant_response = extract_assistant_response(generated_text)
+
+    del input_ids
+    del outputs
+    torch.cuda.empty_cache()
     return assistant_response
 
 def generate_response_7B_instruct(chat_history, max_new_tokens):
@@ -253,18 +286,23 @@ def generate_response_7B_instruct(chat_history, max_new_tokens):
     if input_ids.size(1) > max_input_length:
         input_ids = input_ids[:, -max_input_length:]  # Truncate from the left (oldest tokens)
 
-    outputs = model_7B_instruct.generate(
-        input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
-        return_dict_in_generate=False,
-        output_scores=False,
-        use_cache=True,
-        num_beams=1,
-        do_sample=False
-    )
+    with torch.no_grad():
+        outputs = model_7B_instruct.generate(
+            input_ids=input_ids,
+            max_new_tokens=max_new_tokens,
+            return_dict_in_generate=False,
+            output_scores=False,
+            use_cache=True,
+            num_beams=1,
+            do_sample=False
+        )
 
     generated_text = tokenizer_7B_instruct.decode(outputs[0])
     assistant_response = extract_assistant_response(generated_text)
+
+    del input_ids
+    del outputs
+    torch.cuda.empty_cache()
     return assistant_response
 
 with gr.Blocks() as demo:
@@ -335,3 +373,4 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.queue().launch()
+'''