Files changed (1) hide show
  1. handler.py +5 -4
handler.py CHANGED
@@ -86,18 +86,19 @@ class EndpointHandler:
86
  **inputs,
87
  **gen_params
88
  )
89
- logger.debug(f"Output shape: {output_tokens.shape}")
90
 
91
  # Decode the response
92
- logger.debug("Decoding response")
93
  output_text = self.tokenizer.batch_decode(output_tokens)[0]
94
 
95
  # Extract the assistant's response by removing the input prompt
96
  response = output_text[len(prompt):].strip()
97
  logger.info(f"Generated response length: {len(response)}")
98
- logger.debug(f"Generated response: {response}")
99
 
100
- return [{"role": "assistant", "content": response}]
 
101
 
102
  except Exception as e:
103
  logger.error(f"Error during generation: {str(e)}", exc_info=True)
 
86
  **inputs,
87
  **gen_params
88
  )
89
+ logger.info(f"Output shape: {output_tokens.shape}")
90
 
91
  # Decode the response
92
+ logger.info("Decoding response")
93
  output_text = self.tokenizer.batch_decode(output_tokens)[0]
94
 
95
  # Extract the assistant's response by removing the input prompt
96
  response = output_text[len(prompt):].strip()
97
  logger.info(f"Generated response length: {len(response)}")
98
+ logger.info(f"Generated response: {response}")
99
 
100
+ #return [{"role": "assistant", "content": response}]
101
+ return response
102
 
103
  except Exception as e:
104
  logger.error(f"Error during generation: {str(e)}", exc_info=True)