search_r1/llm_agent/generation.py

• GenerationConfig · class · L15-L24 — class GenerationConfig
• LLMGenerationManager · class · L26-L692 — class LLMGenerationManager
• init · method · L27-L44 — def init( self, tokenizer, actor_rollout_wg, config: GenerationConfig, is_validation: bool = False, )
• _batch_tokenize · method · L46-L53 — def _batch_tokenize(self, responses: List[str]) -> torch.Tensor
• _postprocess_responses · method · L55-L94 — def _postprocess_responses(self, responses: torch.Tensor, is_final: bool = False) -> Tuple[torch.Tensor, List[str]]
• _process_next_obs · method · L96-L109 — def _process_next_obs(self, next_obs: List[str]) -> torch.Tensor
• _update_rolling_state · method · L111-L133 — def _update_rolling_state(self, rollings: DataProto, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor) -> Dict
• _info_masked_concatenate_with_padding · method · L135-L158 — def _info_masked_concatenate_with_padding(self, prompt: torch.Tensor, prompt_with_mask: torch.Tensor, response: torch.Tensor, info: torch.Tensor = None, pad_to_left: bool = True ) -> torch.Tensor
• _update_right_side · method · L160-L182 — def _update_right_side(self, right_side: Dict, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor = None) -> Dict
• _generate_with_gpu_padding · method · L184-L218 — def _generate_with_gpu_padding(self, active_batch: DataProto) -> DataProto
• postprocess_predictions · method · L220-L293 — def postprocess_predictions(self, predictions: List[Any], is_final: bool = False) -> Tuple[List[str], List[str], List[Any], List[Any]]
• execute_predictions · method · L295-L339 — def execute_predictions(self, predictions: List[str], pad_token: str, active_mask=None, do_search=True, is_final=False) -> Tuple[List[str], List[int], List[int], List[int], List[Any], List[Any]]
• batch_search · method · L341-L344 — def batch_search(self, queries: List[str] = None) -> str
• _batch_search · method · L346-L352 — def _batch_search(self, queries)
• _passages2string · method · L354-L361 — def _passages2string(self, retrieval_result)
• run_llm_loop · method · L363-L662 — def run_llm_loop(self, gen_batch, initial_input_ids: torch.Tensor) -> Tuple[Dict, Dict]
• _compose_final_output · method · L664-L692 — def _compose_final_output(self, left_side: Dict, right_side: Dict, meta_info: Dict) -> Tuple[Dict, Dict]