Update utils/models.py

utils/models.py

@@ -154,6 +154,9 @@ def run_inference(model_name, context, question):
         if generation_interrupt.is_set():
             return ""
 
+        # Create interrupt criteria for this generation
+        interrupt_criteria = InterruptCriteria(generation_interrupt)
+
         print("REACHED HERE BEFORE pipe")
         print(f"Loading model {model_name}...")
         if "bitnet" in model_name.lower():
@@ -206,7 +209,10 @@ def run_inference(model_name, context, question):
             result = pipe(
                 text_input,
                 max_new_tokens=512,
-                generation_kwargs={"skip_special_tokens": True},
+                generation_kwargs={
+                    "skip_special_tokens": True,
+                    "stopping_criteria": [interrupt_criteria]  # ADD INTERRUPT SUPPORT
+                },
             )[0]["generated_text"]
 
             result = result[-1]["content"]
@@ -221,7 +227,6 @@ def run_inference(model_name, context, question):
                 **tokenizer_kwargs, 
             )
 
-
             model_inputs = model_inputs.to(model.device)
 
             input_ids = model_inputs.input_ids
@@ -239,7 +244,8 @@ def run_inference(model_name, context, question):
                     attention_mask=attention_mask,
                     max_new_tokens=512,
                     eos_token_id=tokenizer.eos_token_id, 
-                    pad_token_id=tokenizer.pad_token_id  # Addresses the warning
+                    pad_token_id=tokenizer.pad_token_id,
+                    stopping_criteria=[interrupt_criteria]  # ADD INTERRUPT SUPPORT
                 )
             
             generated_token_ids = output_sequences[0][prompt_tokens_length:]
@@ -259,6 +265,7 @@ def run_inference(model_name, context, question):
         #         output_sequences = bitnet_model.generate(
         #             **formatted,
         #             max_new_tokens=512,
+        #             stopping_criteria=[interrupt_criteria]  # ADD INTERRUPT SUPPORT
         #         )
 
         #         result = tokenizer.decode(output_sequences[0][formatted['input_ids'].shape[-1]:], skip_special_tokens=True)
@@ -275,7 +282,10 @@ def run_inference(model_name, context, question):
             outputs = pipe(
                 formatted,
                 max_new_tokens=512,
-                generation_kwargs={"skip_special_tokens": True},
+                generation_kwargs={
+                    "skip_special_tokens": True,
+                    "stopping_criteria": [interrupt_criteria]  # ADD INTERRUPT SUPPORT
+                },
             )
             # print(outputs[0]['generated_text'])
             result = outputs[0]["generated_text"][input_length:]
@@ -290,4 +300,4 @@ def run_inference(model_name, context, question):
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-    return result
+    return result