fix(annottation)

ovuruska · May 7, 2024 · dd1bb95 · dd1bb95
1 parent fe44774
commit dd1bb95
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/deepinfra/models/base/text_generation.py b/deepinfra/models/base/text_generation.py
@@ -4,20 +4,18 @@
 """
 
 import json
-from typing import Union
 
 from deepinfra.models.base import BaseModel
-from deepinfra.types.text_generation.request import TextGenerationRequest
 from deepinfra.types.text_generation.response import TextGenerationResponse
 
 
-class TextGeneration(BaseModel) -> TextGenerationResponse:
+class TextGeneration(BaseModel):
     """
     Initializes one of the DeepInfra text generation models.
     @docs Check the available models at https://deepinfra.com/models/text-generation
     """
 
-    def generate(self, body: dict):
+    def generate(self, body: dict) -> TextGenerationResponse:
         """
         Generates text.
         :param body: