black reformatted files

Author: kumarijy

modules/openvino_code/server/src/app.py

@@ -54,7 +54,9 @@ class GenerationDocStringRequest(BaseModel):
         description="Doc string format passed from extension settings [google | numpy | sphinx | dockblockr | ...]",
         example="numpy",
     )
-    definition: str = Field("", description="Function signature", example="def fibonacci(n):")
+    definition: str = Field(
+        "", description="Function signature", example="def fibonacci(n):"
+    )
     parameters: GenerationParameters
 
 
@@ -111,10 +113,16 @@ async def generate_stream(
     request: Request,
     generator: GeneratorFunctor = Depends(get_generator_dummy),
 ) -> StreamingResponse:
-    generation_request = TypeAdapter(GenerationRequest).validate_python(await request.json())
+    generation_request = TypeAdapter(GenerationRequest).validate_python(
+        await request.json()
+    )
     logger.info(generation_request)
     return StreamingResponse(
-        generator.generate_stream(generation_request.inputs, generation_request.parameters.model_dump(), request)
+        generator.generate_stream(
+            generation_request.inputs,
+            generation_request.parameters.model_dump(),
+            request,
+        )
     )
 
 
@@ -127,7 +135,11 @@ async def summarize(
 
     start = perf_counter()
     generated_text: str = generator.summarize(
-        request.inputs, request.template, request.definition, request.format, request.parameters.model_dump()
+        request.inputs,
+        request.template,
+        request.definition,
+        request.format,
+        request.parameters.model_dump(),
     )
     stop = perf_counter()
 
@@ -148,6 +160,10 @@ async def summarize_stream(
     logger.info(request)
     return StreamingResponse(
         generator.summarize_stream(
-            request.inputs, request.template, request.definition, request.format, request.parameters.model_dump()
+            request.inputs,
+            request.template,
+            request.definition,
+            request.format,
+            request.parameters.model_dump(),
         )
     )

modules/openvino_code/server/src/generators.py

@@ -5,7 +5,17 @@
 from pathlib import Path
 from threading import Thread
 from time import time
-from typing import Any, Callable, Container, Dict, Generator, List, Optional, Type, Union
+from typing import (
+    Any,
+    Callable,
+    Container,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Type,
+    Union,
+)
 
 import torch
 from fastapi import Request
@@ -30,14 +40,18 @@
 model_dir = Path("models")
 model_dir.mkdir(exist_ok=True)
 
-SUMMARIZE_INSTRUCTION = "{function}\n\n# The function with {style} style docstring\n\n{signature}\n"
+SUMMARIZE_INSTRUCTION = (
+    "{function}\n\n# The function with {style} style docstring\n\n{signature}\n"
+)
 SUMMARIZE_STOP_TOKENS = ("\r\n", "\n")
 
 
 def get_model_class(checkpoint: Union[str, Path]) -> Type[OVModel]:
     config = AutoConfig.from_pretrained(checkpoint)
     architecture: str = config.architectures[0]
-    if architecture.endswith("ConditionalGeneration") or architecture.endswith("Seq2SeqLM"):
+    if architecture.endswith("ConditionalGeneration") or architecture.endswith(
+        "Seq2SeqLM"
+    ):
         return OVModelForSeq2SeqLM
 
     return OVModelForCausalLM
@@ -48,16 +62,27 @@ def get_model(checkpoint: str, device: str = "CPU") -> OVModel:
     model_path = model_dir / Path(checkpoint)
     if model_path.exists():
         model_class = get_model_class(model_path)
-        model = model_class.from_pretrained(model_path, ov_config=ov_config, compile=False, device=device)
+        model = model_class.from_pretrained(
+            model_path, ov_config=ov_config, compile=False, device=device
+        )
     else:
         model_class = get_model_class(checkpoint)
         try:
             model = model_class.from_pretrained(
-                checkpoint, ov_config=ov_config, compile=False, device=device, trust_remote_code=True
+                checkpoint,
+                ov_config=ov_config,
+                compile=False,
+                device=device,
+                trust_remote_code=True,
             )
         except EntryNotFoundError:
             model = model_class.from_pretrained(
-                checkpoint, ov_config=ov_config, export=True, compile=False, device=device, trust_remote_code=True
+                checkpoint,
+                ov_config=ov_config,
+                export=True,
+                compile=False,
+                device=device,
+                trust_remote_code=True,
             )
         model.save_pretrained(model_path)
     model.compile()
@@ -72,13 +97,29 @@ class GeneratorFunctor:
     def __call__(self, input_text: str, parameters: Dict[str, Any]) -> str:
         raise NotImplementedError
 
-    async def generate_stream(self, input_text: str, parameters: Dict[str, Any], request: Request):
+    async def generate_stream(
+        self, input_text: str, parameters: Dict[str, Any], request: Request
+    ):
         raise NotImplementedError
 
-    def summarize(self, input_text: str, template: str, signature: str, style: str, parameters: Dict[str, Any]):
+    def summarize(
+        self,
+        input_text: str,
+        template: str,
+        signature: str,
+        style: str,
+        parameters: Dict[str, Any],
+    ):
         raise NotImplementedError
 
-    def summarize_stream(self, input_text: str, template: str, signature: str, style: str, parameters: Dict[str, Any]):
+    def summarize_stream(
+        self,
+        input_text: str,
+        template: str,
+        signature: str,
+        style: str,
+        parameters: Dict[str, Any],
+    ):
         raise NotImplementedError
 
 
@@ -113,9 +154,14 @@ def __init__(
         if summarize_stop_tokens:
             stop_tokens = []
             for token_id in self.tokenizer.vocab.values():
-                if any(stop_word in self.tokenizer.decode(token_id) for stop_word in summarize_stop_tokens):
+                if any(
+                    stop_word in self.tokenizer.decode(token_id)
+                    for stop_word in summarize_stop_tokens
+                ):
                     stop_tokens.append(token_id)
-            self.summarize_stopping_criteria = StoppingCriteriaList([StopOnTokens(stop_tokens)])
+            self.summarize_stopping_criteria = StoppingCriteriaList(
+                [StopOnTokens(stop_tokens)]
+            )
 
     def __call__(self, input_text: str, parameters: Dict[str, Any]) -> str:
         input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
@@ -126,20 +172,36 @@ def __call__(self, input_text: str, parameters: Dict[str, Any]) -> str:
             stopping_criteria = StoppingCriteriaList([stop_on_time])
 
         prompt_len = input_ids.shape[-1]
-        config = GenerationConfig.from_dict({**self.generation_config.to_dict(), **parameters})
+        config = GenerationConfig.from_dict(
+            {**self.generation_config.to_dict(), **parameters}
+        )
         output_ids = self.model.generate(
-            input_ids, generation_config=config, stopping_criteria=stopping_criteria, **self.assistant_model_config
+            input_ids,
+            generation_config=config,
+            stopping_criteria=stopping_criteria,
+            **self.assistant_model_config,
         )[0][prompt_len:]
-        logger.info(f"Number of input tokens: {prompt_len}; generated {len(output_ids)} tokens")
-        return self.tokenizer.decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        logger.info(
+            f"Number of input tokens: {prompt_len}; generated {len(output_ids)} tokens"
+        )
+        return self.tokenizer.decode(
+            output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
 
     async def generate_stream(
-        self, input_text: str, parameters: Dict[str, Any], request: Optional[Request] = None
+        self,
+        input_text: str,
+        parameters: Dict[str, Any],
+        request: Optional[Request] = None,
     ) -> Generator[str, None, None]:
         input_ids = self.tokenizer.encode(input_text, return_tensors="pt")
-        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+        streamer = TextIteratorStreamer(
+            self.tokenizer, skip_prompt=True, skip_special_tokens=True
+        )
         parameters["streamer"] = streamer
-        config = GenerationConfig.from_dict({**self.generation_config.to_dict(), **parameters})
+        config = GenerationConfig.from_dict(
+            {**self.generation_config.to_dict(), **parameters}
+        )
 
         stop_on_tokens = StopOnTokens([])
 
@@ -180,7 +242,9 @@ def generate_between(
         parameters: Dict[str, Any],
         stopping_criteria: Optional[StoppingCriteriaList] = None,
     ) -> str:
-        config = GenerationConfig.from_dict({**self.generation_config.to_dict(), **parameters})
+        config = GenerationConfig.from_dict(
+            {**self.generation_config.to_dict(), **parameters}
+        )
 
         prompt = torch.tensor([[]], dtype=torch.int64)
         buffer = StringIO()
@@ -192,13 +256,20 @@ def generate_between(
             prev_len = prompt.shape[-1]
 
             prompt = self.model.generate(
-                prompt, generation_config=config, stopping_criteria=stopping_criteria, **self.assistant_model_config
+                prompt,
+                generation_config=config,
+                stopping_criteria=stopping_criteria,
+                **self.assistant_model_config,
             )[
                 :, :-1
             ]  # skip the last token - stop token
 
-            decoded = self.tokenizer.decode(prompt[0, prev_len:], skip_special_tokens=True)
-            buffer.write(decoded.lstrip(" "))  # hack to delete leadding spaces if there are any
+            decoded = self.tokenizer.decode(
+                prompt[0, prev_len:], skip_special_tokens=True
+            )
+            buffer.write(
+                decoded.lstrip(" ")
+            )  # hack to delete leadding spaces if there are any
         buffer.write(input_parts[-1])
         return buffer.getvalue()
 
@@ -208,7 +279,9 @@ async def generate_between_stream(
         parameters: Dict[str, Any],
         stopping_criteria: Optional[StoppingCriteriaList] = None,
     ) -> Generator[str, None, None]:
-        config = GenerationConfig.from_dict({**self.generation_config.to_dict(), **parameters})
+        config = GenerationConfig.from_dict(
+            {**self.generation_config.to_dict(), **parameters}
+        )
 
         prompt = self.tokenizer.encode(input_parts[0], return_tensors="pt")
         for text_input in input_parts[1:-1]:
@@ -219,12 +292,17 @@ async def generate_between_stream(
             prev_len = prompt.shape[-1]
 
             prompt = self.model.generate(
-                prompt, generation_config=config, stopping_criteria=stopping_criteria, **self.assistant_model_config
+                prompt,
+                generation_config=config,
+                stopping_criteria=stopping_criteria,
+                **self.assistant_model_config,
             )[
                 :, :-1
             ]  # skip the last token - stop token
 
-            decoded = self.tokenizer.decode(prompt[0, prev_len:], skip_special_tokens=True)
+            decoded = self.tokenizer.decode(
+                prompt[0, prev_len:], skip_special_tokens=True
+            )
             yield decoded.lstrip(" ")  # hack to delete leadding spaces if there are any
 
         yield input_parts[-1]
@@ -237,24 +315,40 @@ def summarization_input(function: str, signature: str, style: str) -> str:
             signature=signature,
         )
 
-    def summarize(self, input_text: str, template: str, signature: str, style: str, parameters: Dict[str, Any]) -> str:
+    def summarize(
+        self,
+        input_text: str,
+        template: str,
+        signature: str,
+        style: str,
+        parameters: Dict[str, Any],
+    ) -> str:
         prompt = self.summarization_input(input_text, signature, style)
         splited_template = re.split(r"\$\{.*\}", template)
         splited_template[0] = prompt + splited_template[0]
 
-        return self.generate_between(splited_template, parameters, stopping_criteria=self.summarize_stopping_criteria)[
-            len(prompt) :
-        ]
+        return self.generate_between(
+            splited_template,
+            parameters,
+            stopping_criteria=self.summarize_stopping_criteria,
+        )[len(prompt) :]
 
     async def summarize_stream(
-        self, input_text: str, template: str, signature: str, style: str, parameters: Dict[str, Any]
+        self,
+        input_text: str,
+        template: str,
+        signature: str,
+        style: str,
+        parameters: Dict[str, Any],
     ):
         prompt = self.summarization_input(input_text, signature, style)
         splited_template = re.split(r"\$\{.*\}", template)
         splited_template = [prompt] + splited_template
 
         async for token in self.generate_between_stream(
-            splited_template, parameters, stopping_criteria=self.summarize_stopping_criteria
+            splited_template,
+            parameters,
+            stopping_criteria=self.summarize_stopping_criteria,
         ):
             yield token
 
@@ -279,7 +373,9 @@ def __init__(self, token_ids: List[int]) -> None:
         self.cancelled = False
         self.token_ids = torch.tensor(token_ids, requires_grad=False)
 
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
+    ) -> bool:
         if self.cancelled:
             return True
         return torch.any(torch.eq(input_ids[0, -1], self.token_ids)).item()
@@ -304,4 +400,6 @@ def __call__(self, *args, **kwargs) -> bool:
         self.time_for_prev_token = elapsed
         self.time = current_time
 
-        return self.stop_until < current_time + self.time_for_prev_token * self.grow_factor
+        return (
+            self.stop_until < current_time + self.time_for_prev_token * self.grow_factor
+        )

modules/openvino_code/server/src/utils.py

@@ -9,7 +9,9 @@ def get_parser() -> argparse.ArgumentParser:
     parser.add_argument("--port", type=int, default="8000")
 
     parser.add_argument("--model", type=str, required=True)
-    parser.add_argument("--tokenizer_checkpoint", type=str, required=False, default=None)
+    parser.add_argument(
+        "--tokenizer_checkpoint", type=str, required=False, default=None
+    )
     parser.add_argument("--device", type=str, required=False, default="CPU")
     parser.add_argument("--assistant", type=str, required=False, default=None)
     parser.add_argument("--summarization-endpoint", action="store_true")
@@ -34,7 +36,9 @@ def get_logger(
 class ServerLogger(logging.Logger):
     _server_log_prefix = "[OpenVINO Code Server Log]"
 
-    default_formatter = logging.Formatter(f"{_server_log_prefix} %(asctime)s %(levelname)s %(message)s")
+    default_formatter = logging.Formatter(
+        f"{_server_log_prefix} %(asctime)s %(levelname)s %(message)s"
+    )
 
     def __init__(self, name):
         super(ServerLogger, self).__init__(name)