Add types to _TextGenerationTimings attributes

Revert to using timer.time for `TOKEN_GENERATION`
Remove finally clause from `contextmanagers`
Address review comments from @rahul-tuli

Author: rahul-tuli

src/deepsparse/transformers/pipelines/text_generation.py

@@ -35,10 +35,10 @@
 
 @dataclass(frozen=True)
 class _TextGenerationTimings:
-    PROMPT_PREFILL = "engine_prompt_prefill"
-    PROMPT_PREFILL_SINGLE = "engine_prompt_prefill_single"
-    TOKEN_GENERATION = "engine_token_generation"
-    TOKEN_GENERATION_SINGLE = "engine_token_generation_single"
+    PROMPT_PREFILL: str = "engine_prompt_prefill"
+    PROMPT_PREFILL_SINGLE: str = "engine_prompt_prefill_single"
+    TOKEN_GENERATION: str = "engine_token_generation"
+    TOKEN_GENERATION_SINGLE: str = "engine_token_generation_single"
 
 
 class TextGenerationInput(BaseModel):
@@ -344,17 +344,19 @@ def engine_forward(
             generated_tokens = [tokens[-1]]
             generated_logits = prompt_logits
 
-            timer.start(_TextGenerationTimings.TOKEN_GENERATION)
-            while len(generated_tokens) < max_tokens:
-                with timer.time(_TextGenerationTimings.TOKEN_GENERATION_SINGLE):
-                    token, logits = self.autoregressive_inference(tokens)
-                tokens.append(token)
-                generated_tokens.append(token)
-                generated_logits.append(logits)
-
-                if token == self.tokenizer.eos_token_id and not self.force_max_tokens:
-                    break
-            timer.stop(_TextGenerationTimings.TOKEN_GENERATION)
+            with timer.time(_TextGenerationTimings.TOKEN_GENERATION):
+                while len(generated_tokens) < max_tokens:
+                    with timer.time(_TextGenerationTimings.TOKEN_GENERATION_SINGLE):
+                        token, logits = self.autoregressive_inference(tokens)
+                    tokens.append(token)
+                    generated_tokens.append(token)
+                    generated_logits.append(logits)
+
+                    if (
+                        token == self.tokenizer.eos_token_id
+                        and not self.force_max_tokens
+                    ):
+                        break
 
         return numpy.array([generated_tokens]), numpy.concatenate(
             generated_logits, axis=1

src/deepsparse/utils/timer.py

@@ -128,11 +128,8 @@ def time(self, stage: str):
         :param stage: the name of the stage to time
         """
         self.start(stage)
-
-        try:
-            yield
-        finally:
-            self.stop(stage)
+        yield
+        self.stop(stage)
 
     def start(self, stage: str):
         """
@@ -363,9 +360,6 @@ def new_timer_context(self, total_inference: bool = True) -> StagedTimer:
             self._timers = [timer]
 
         timer_context.set(timer)
-
-        try:
-            yield timer
-        finally:
-            if total_inference:
-                timer.stop(InferenceStages.TOTAL_INFERENCE)
+        yield timer
+        if total_inference:
+            timer.stop(InferenceStages.TOTAL_INFERENCE)