cleanup job queue on failure

outlines/models/exllamav2.py

@@ -1,6 +1,7 @@
 import dataclasses
 from typing import TYPE_CHECKING, Iterator, List, Optional, Tuple, TypedDict, Union
 
+import torch
 from typing_extensions import Unpack
 
 from outlines.generate.api import GenerationParameters, SamplingParameters
@@ -29,7 +30,9 @@ def convert_token_to_string(self, token):
         return token
 
     def decode(self, token_ids: "torch.LongTensor") -> List[str]:
-        return self.exl2_tokenizer.batch_decode(token_ids, decode_special_tokens=False)
+        return self.exl2_tokenizer.decode(
+            torch.tensor([token_ids]), decode_special_tokens=False
+        )
 
 
 class ExLlamaV2Model:
@@ -219,16 +222,21 @@ def stream(
         next_text = [""] * batch_size
 
         def token_generator() -> Iterator[str]:
-            while self.generator.num_remaining_jobs():
-                results = self.generator.iterate()
-                for r in results:
-                    idx = order[r["serial"]]
-                    if r["stage"] == "streaming":
-                        text = r.get("text", "")
-                        next_text[idx] = text
-                    if r["eos"]:
-                        next_text[idx] = ""
-                yield self.reformat_output(next_text, sampling_parameters)
+            try:
+                while self.generator.num_remaining_jobs():
+                    results = self.generator.iterate()
+                    for r in results:
+                        idx = order[r["serial"]]
+                        if r["stage"] == "streaming":
+                            text = r.get("text", "")
+                            next_text[idx] = text
+                        if r["eos"]:
+                            next_text[idx] = ""
+                    yield self.reformat_output(next_text, sampling_parameters)
+            except Exception as e:
+                for job in self.generator.pending_jobs:
+                    job.cancel()
+                raise e
             return
 
         return token_generator()