Update src/txagent/txagent.py

src/txagent/txagent.py

@@ -80,11 +80,15 @@ class TxAgent:
         if model_name:
             self.model_name = model_name
 
-        self.model = LLM(model=self.model_name, dtype="float16")
-        self.chat_template = Template(self.model.get_tokenizer().chat_template)
-        self.tokenizer = self.model.get_tokenizer()
-        logger.info("Model %s loaded successfully", self.model_name)
-        return f"Model {self.model_name} loaded successfully."
+        try:
+            self.model = LLM(model=self.model_name, dtype="float16", max_model_len=131072)
+            self.chat_template = Template(self.model.get_tokenizer().chat_template)
+            self.tokenizer = self.model.get_tokenizer()
+            logger.info("Model %s loaded successfully", self.model_name)
+            return f"Model {self.model_name} loaded successfully."
+        except Exception as e:
+            logger.error(f"Model loading error: {e}")
+            raise
 
     def load_tooluniverse(self):
         self.tooluniverse = ToolUniverse(tool_files=self.tool_files_dict)
@@ -416,17 +420,23 @@ class TxAgent:
             prompt += output_begin_string
 
         if check_token_status and max_token:
-            num_input_tokens = len(self.tokenizer.encode(prompt, return_tensors="pt")[0])
+            num_input_tokens = len(self.tokenizer.encode(prompt, add_special_tokens=False))
             if num_input_tokens > max_token:
-                torch.cuda.empty_cache()
-                gc.collect()
-                logger.info("Token overflow: %d > %d", num_input_tokens, max_token)
+                logger.warning(f"Input tokens ({num_input_tokens}) exceed max_token ({max_token}). Truncating.")
+                prompt_tokens = self.tokenizer.encode(prompt, add_special_tokens=False)[:max_token]
+                prompt = self.tokenizer.decode(prompt_tokens)
+            if num_input_tokens > 131072:
+                logger.error(f"Input tokens ({num_input_tokens}) exceed model limit (131072).")
                 return None, True
-            logger.debug("Input tokens: %d", num_input_tokens)
 
-        output = model.generate(prompt, sampling_params=sampling_params)
-        output = output[0].outputs[0].text
-        logger.debug("Inference output: %s", output[:100])
+        try:
+            output = model.generate(prompt, sampling_params=sampling_params)
+            output = output[0].outputs[0].text
+            logger.debug("Inference output: %s", output[:100])
+        except Exception as e:
+            logger.error(f"Inference error: {e}")
+            return None, True
+
         torch.cuda.empty_cache()
         if check_token_status:
             return output, False
@@ -442,17 +452,20 @@ Patient Record Excerpt:
 """
         conversation = self.set_system_prompt([], prompt.format(chunk=message))
         conversation.append({"role": "user", "content": message})
-        output = self.llm_infer(
+        output, token_overflow = self.llm_infer(
             messages=conversation,
             temperature=temperature,
             max_new_tokens=max_new_tokens,
             max_token=max_token,
             tools=[]  # No tools
         )
-        if '[FinalAnswer]' in output:
+        if token_overflow:
+            logger.error("Token overflow in quick summary")
+            return "Error: Input too large for quick summary."
+        if output and '[FinalAnswer]' in output:
             output = output.split('[FinalAnswer]')[-1].strip()
-        logger.debug("Quick summary output: %s", output[:100])
-        return output
+        logger.debug("Quick summary output: %s", output[:100] if output else "None")
+        return output or "No missed diagnoses identified"
 
     def run_background_report(self, message: str, history: list, temperature: float,
                              max_new_tokens: int, max_token: int, call_agent: bool,
@@ -539,10 +552,10 @@ Patient Record Excerpt:
                     f.write(combined_response)
                 logger.info("Detailed report saved to %s", report_path)
             except Exception as e:
-                logger.error("Failed to save report: %s", e)
+                logger.error(f"Failed to save report: {e}")
 
         except Exception as e:
-            logger.error("Background report error: %s", e)
+            logger.error(f"Background report error: {e}")
             combined_response += f"Error: {e}\n"
             with open(report_path, "w", encoding="utf-8") as f:
                 f.write(combined_response)