fix: delete useless prompt_logprobs=1

Author: ChenZiHong-Gavin

graphgen/bases/base_llm_wrapper.py

@@ -26,11 +26,11 @@ def __init__(
         **kwargs: Any,
     ):
         self.system_prompt = system_prompt
-        self.temperature = temperature
-        self.max_tokens = max_tokens
-        self.repetition_penalty = repetition_penalty
-        self.top_p = top_p
-        self.top_k = top_k
+        self.temperature = float(temperature)
+        self.max_tokens = int(max_tokens)
+        self.repetition_penalty = float(repetition_penalty)
+        self.top_p = float(top_p)
+        self.top_k = int(top_k)
         self.tokenizer = tokenizer
 
         for k, v in kwargs.items():

graphgen/models/llm/local/vllm_wrapper.py

@@ -20,13 +20,9 @@ def __init__(
         temperature: float = 0.6,
         top_p: float = 1.0,
         top_k: int = 5,
-        timeout: float = 300
+        timeout: float = 300,
         **kwargs: Any,
     ):
-        temperature = float(temperature)
-        top_p = float(top_p)
-        top_k = int(top_k)
-
         super().__init__(temperature=temperature, top_p=top_p, top_k=top_k, **kwargs)
         try:
             from vllm import AsyncEngineArgs, AsyncLLMEngine, SamplingParams
@@ -45,10 +41,7 @@ def __init__(
             disable_log_stats=False,
         )
         self.engine = AsyncLLMEngine.from_engine_args(engine_args)
-        self.temperature = temperature
-        self.top_p = top_p
-        self.top_k = top_k
-        self.timeout = timeout
+        self.timeout = float(timeout)
 
     @staticmethod
     def _build_inputs(prompt: str, history: Optional[List[str]] = None) -> str:
@@ -89,20 +82,15 @@ async def generate_answer(
                 self._consume_generator(result_generator),
                 timeout=self.timeout
             )
-            
+
             if not final_output or not final_output.outputs:
                 return ""
 
             result_text = final_output.outputs[0].text
             return result_text
-            
-        except asyncio.TimeoutError:
-            await self.engine.abort(request_id)
-            raise
-        except asyncio.CancelledError:
-            await self.engine.abort(request_id)
-            raise
+
         except Exception as e:
+            print(f"Error in generate_answer: {e}")
             await self.engine.abort(request_id)
             raise
 
@@ -116,7 +104,6 @@ async def generate_topk_per_token(
             temperature=0,
             max_tokens=1,
             logprobs=self.top_k,
-            prompt_logprobs=1,
         )
 
         result_generator = self.engine.generate(full_prompt, sp, request_id=request_id)
@@ -126,7 +113,7 @@ async def generate_topk_per_token(
                 self._consume_generator(result_generator),
                 timeout=self.timeout
             )
-            
+
             if (
                 not final_output
                 or not final_output.outputs
@@ -154,14 +141,9 @@ async def generate_topk_per_token(
                 )
                 return [main_token]
             return []
-            
-        except asyncio.TimeoutError:
-            await self.engine.abort(request_id)
-            raise
-        except asyncio.CancelledError:
-            await self.engine.abort(request_id)
-            raise
+
         except Exception as e:
+            print(f"Error in generate_topk_per_token: {e}")
             await self.engine.abort(request_id)
             raise
 
@@ -171,4 +153,3 @@ async def generate_inputs_prob(
         raise NotImplementedError(
             "VLLMWrapper does not support per-token logprobs yet."
         )
-