Fixing dockerfile v5

src/api.py

@@ -27,16 +27,17 @@ class InferenceApi(LitAPI):
         )
         self.logger.info(f"Inference API setup completed on device: {device}")
 
-    async def predict(self, x: str, **kwargs) -> Union[str, Iterator[str]]:
-        """
-        Main prediction method required by LitAPI.
-        If streaming is enabled, yields chunks; otherwise returns complete response.
-        """
-        if self.stream:
-            async for chunk in self.generate_stream(x, **kwargs):
-                yield chunk
-        else:
-            return await self.generate_response(x, **kwargs)
+async def predict(self, x: str, **kwargs) -> Union[str, Iterator[str]]:
+    """
+    Main prediction method required by LitAPI.
+    If streaming is enabled, yields chunks; otherwise returns via yield.
+    """
+    if self.stream:
+        async for chunk in self.generate_stream(x, **kwargs):
+            yield chunk
+    else:
+        response = await self.generate_response(x, **kwargs)
+        yield response
 
     def decode_request(self, request: Any, **kwargs) -> str:
         """Convert the request payload to input format."""