MCPcopy Create free account
hub / github.com/agentscope-ai/Trinity-RFT / test_generate

Method test_generate

tests/common/vllm_test.py:181–346  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

179 self.assertGreater(len(choice.routed_experts), 0)
180
181 async def test_generate(self): # noqa: C901
182 self.assertEqual(self.model_wrapper.model_path, self.config.model.model_path)
183 prompts = ["Hello, world!", "Hello, my name is"]
184 n = self.config.algorithm.repeat_times
185 if self.use_async:
186 generate_results = await self.model_wrapper.generate_async(
187 prompts, n=n, temperature=1.0
188 )
189 else:
190 generate_results = self.model_wrapper.generate(prompts, n=n, temperature=1.0)
191 self.assertEqual(len(generate_results), len(prompts) * n)
192 if self.enable_return_routed_experts:
193 for exp in generate_results:
194 _assert_routed_experts_shape(
195 self,
196 exp,
197 self.expected_routed_experts_layers,
198 self.expected_routed_experts_topk,
199 )
200 if self.config.explorer.rollout_model.enable_history:
201 history_experiences = self.model_wrapper.extract_experience_from_history(
202 clear_history=False
203 )
204 self.assertEqual(len(history_experiences), len(generate_results))
205 for exp, history_exp in zip(generate_results, history_experiences):
206 self.assertEqual(exp.response_text, history_exp.response_text)
207 self.assertEqual(exp.tokens.tolist(), history_exp.tokens.tolist())
208 self.assertEqual(exp.prompt_length, history_exp.prompt_length)
209 self.assertEqual(exp.logprobs.tolist(), history_exp.logprobs.tolist())
210 if self.enable_return_routed_experts:
211 _assert_routed_experts_shape(
212 self,
213 history_exp,
214 self.expected_routed_experts_layers,
215 self.expected_routed_experts_topk,
216 )
217 else:
218 with self.assertRaises(ValueError):
219 self.model_wrapper.extract_experience_from_history(clear_history=False)
220 messages = [
221 {"role": "system", "content": "You are a helpful assistant."},
222 {"role": "user", "content": "What's the weather like today?"},
223 {
224 "role": "assistant",
225 "content": "I'm sorry, but as an AI language model, I don't have access to real-time weather information. To get accurate weather information for your location, you can check a weather website or app, or look outside if possible.",
226 },
227 {"role": "user", "content": "OK, thanks!"},
228 ]
229 if self.use_async:
230 results = await self.model_wrapper.chat_async(messages, n=n, temperature=1.0)
231 else:
232 results = self.model_wrapper.chat(messages, n=n, temperature=1.0)
233 self.assertEqual(len(results), n)
234 if self.enable_return_routed_experts:
235 for exp in results:
236 _assert_routed_experts_shape(
237 self,
238 exp,

Callers

nothing calls this directly

Calls 15

generate_asyncMethod · 0.80
chat_asyncMethod · 0.80
logprobs_asyncMethod · 0.80
apply_chat_templateMethod · 0.80
createMethod · 0.80
generateMethod · 0.45
chatMethod · 0.45
logprobsMethod · 0.45

Tested by

no test coverage detected