diff --git a/deepinfra/models/base/text_generation.py b/deepinfra/models/base/text_generation.py
index 68ad86b..670ba9b 100644
--- 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: