fix(models): fix lint errors

Author: ChenZiHong-Gavin

models/__init__.py

@@ -38,4 +38,4 @@
     "UniEvaluator",
     # strategy models
     "TraverseStrategy",
-]
+]

models/embed/embedding.py

@@ -1,8 +1,7 @@
+from dataclasses import dataclass
 import asyncio
 import numpy as np
 
-from dataclasses import dataclass
-
 class UnlimitedSemaphore:
     """A context manager that allows unlimited access."""
 

models/evaluate/length_evaluator.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from .base_evaluator import BaseEvaluator
+from models.evaluate.base_evaluator import BaseEvaluator
 from models.llm.tokenizer import Tokenizer
 from models.text.text_pair import TextPair
 from utils import create_event_loop
@@ -16,7 +16,7 @@ def __post_init__(self):
     async def evaluate_single(self, pair: TextPair) -> float:
         loop = create_event_loop()
         return await loop.run_in_executor(None, self._calculate_length, pair.answer)
-    
+
     def _calculate_length(self, text: str) -> float:
         tokens = self.tokenizer.encode_string(text)
         return len(tokens)

models/llm/openai_model.py

@@ -1,16 +1,19 @@
 import math
 from dataclasses import dataclass
 from typing import List, Dict, Optional
-from openai import AsyncOpenAI, RateLimitError, APIConnectionError, APITimeoutError, ChatCompletion
-from models import TopkTokenModel, Token
+import openai
+from openai import AsyncOpenAI, RateLimitError, APIConnectionError, APITimeoutError
 from tenacity import (
     retry,
     stop_after_attempt,
     wait_exponential,
     retry_if_exception_type,
 )
 
-def get_top_response_tokens(response: ChatCompletion) -> List[Token]:
+from models import TopkTokenModel, Token
+
+
+def get_top_response_tokens(response: openai.ChatCompletion) -> List[Token]:
     token_logprobs = response.choices[0].logprobs.content
     tokens = []
     for token_prob in token_logprobs:
@@ -76,6 +79,7 @@ async def generate_topk_per_token(self, text: str, history: Optional[List[str]]
 
         completion = await self.client.chat.completions.create(
             model=self.model_name,
+            messages=kwargs["messages"],
             **kwargs
         )
 
@@ -94,7 +98,11 @@ async def generate_answer(self, text: str, history: Optional[List[str]] = None,
 
         completion = await self.client.chat.completions.create(
             model=self.model_name,
+            messages=kwargs["messages"],
             **kwargs
         )
 
         return completion.choices[0].message.content
+
+    async def generate_inputs_prob(self, text: str, history: Optional[List[str]] = None) -> List[Token]:
+        raise NotImplementedError

models/storage/base_storage.py

@@ -11,11 +11,10 @@ class StorageNameSpace:
 
     async def index_done_callback(self):
         """commit the storage operations after indexing"""
-        pass
 
     async def query_done_callback(self):
         """commit the storage operations after querying"""
-        pass
+
 
 @dataclass
 class BaseKVStorage(Generic[T], StorageNameSpace):

models/storage/json_storage.py

@@ -36,7 +36,7 @@ async def get_by_ids(self, ids, fields=None):
         ]
 
     async def filter_keys(self, data: list[str]) -> set[str]:
-        return set([s for s in data if s not in self._data])
+        return {s for s in data if s not in self._data}
 
     async def upsert(self, data: dict):
         left_data = {k: v for k, v in data.items() if k not in self._data}

models/storage/networkx_storage.py

@@ -17,9 +17,7 @@ def load_nx_graph(file_name) -> Optional[nx.Graph]:
 
     @staticmethod
     def write_nx_graph(graph: nx.Graph, file_name):
-        logger.info(
-            f"Writing graph with {graph.number_of_nodes()} nodes, {graph.number_of_edges()} edges"
-        )
+        logger.info("Writing graph with %d nodes, %d edges", graph.number_of_nodes(), graph.number_of_edges())
         nx.write_graphml(graph, file_name)
 
     @staticmethod
@@ -56,9 +54,7 @@ def _stabilize_graph(graph: nx.Graph) -> nx.Graph:
             def _sort_source_target(edge):
                 source, target, edge_data = edge
                 if source > target:
-                    temp = source
-                    source = target
-                    target = temp
+                    source, target = target, source
                 return source, target, edge_data
 
             edges = [_sort_source_target(edge) for edge in edges]
@@ -81,7 +77,8 @@ def __post_init__(self):
         preloaded_graph = NetworkXStorage.load_nx_graph(self._graphml_xml_file)
         if preloaded_graph is not None:
             logger.info(
-                f"Loaded graph from {self._graphml_xml_file} with {preloaded_graph.number_of_nodes()} nodes, {preloaded_graph.number_of_edges()} edges"
+                "Loaded graph from %s with %d nodes, %d edges", self._graphml_xml_file,
+                preloaded_graph.number_of_nodes(), preloaded_graph.number_of_edges()
             )
         self._graph = preloaded_graph or nx.Graph()