| 1829 | self.model_wrapper = self.engines[0] |
| 1830 | |
| 1831 | async def test_generate(self): |
| 1832 | base_dir = os.path.dirname(__file__) |
| 1833 | target_dir = os.path.join(base_dir, "..", "..", "trinity", "trainer", "verl_legacy") |
| 1834 | with open(os.path.join(target_dir, "fsdp_workers.py")) as f: |
| 1835 | fsdp_code = f.read() |
| 1836 | with open(os.path.join(target_dir, "megatron_workers.py")) as f: |
| 1837 | megatron_code = f.read() |
| 1838 | target_dir = os.path.join(base_dir, "..", "..", "trinity", "common") |
| 1839 | with open(os.path.join(target_dir, "config.py")) as f: |
| 1840 | config_code = f.read() |
| 1841 | target_dir = os.path.join(base_dir, "..", "..", "trinity", "manager") |
| 1842 | with open(os.path.join(target_dir, "config_manager.py")) as f: |
| 1843 | config_manager_code = f.read() |
| 1844 | |
| 1845 | messages = [ |
| 1846 | {"role": "system", "content": "You are a helpful assistant."}, |
| 1847 | { |
| 1848 | "role": "user", |
| 1849 | "content": """# Please add comments and documentation for these following code, """ |
| 1850 | """make sure the code is well-structured and easy to read, """ |
| 1851 | """and the complete code must be shown, do not omit any parts.\n""" |
| 1852 | f"""## fsdp_workers.py\n{fsdp_code}\n""" |
| 1853 | f"""## megatron_workers.py\n{megatron_code}\n""" |
| 1854 | f"""## config.py\n{config_code}\n""" |
| 1855 | f"""## config_manager.py\n{config_manager_code}\n""", |
| 1856 | }, |
| 1857 | ] |
| 1858 | response = self.model_wrapper.chat(messages, n=1, temperature=0.7, logprobs=True)[0] |
| 1859 | self.assertGreater( |
| 1860 | response.prompt_length, 40960 |
| 1861 | ) # If not long enough, please add more files to prompt |
| 1862 | self.assertGreater(response.logprobs.shape[0], 1000) |
| 1863 | |
| 1864 | |
| 1865 | class TestTinkerAPI(VLLMTestBase): |