1
0

conftest.py 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791
  1. import contextlib
  2. import gc
  3. import json
  4. import os
  5. import sys
  6. import tempfile
  7. from collections import UserList
  8. from enum import Enum
  9. from typing import (Any, Callable, Dict, List, Optional, Tuple, TypedDict,
  10. TypeVar, Union)
  11. import numpy as np
  12. import pytest
  13. import torch
  14. import torch.nn as nn
  15. import torch.nn.functional as F
  16. from huggingface_hub import snapshot_download
  17. from loguru import logger
  18. from PIL import Image
  19. from transformers import (AutoModelForCausalLM, AutoTokenizer, BatchEncoding,
  20. BatchFeature)
  21. from aphrodite import LLM, SamplingParams
  22. from aphrodite.assets.image import ImageAsset
  23. from aphrodite.common.config import TokenizerPoolConfig
  24. from aphrodite.common.outputs import RequestOutput
  25. from aphrodite.common.sequence import SampleLogprobs
  26. from aphrodite.common.utils import (STR_DTYPE_TO_TORCH_DTYPE,
  27. cuda_device_count_stateless, identity,
  28. is_cpu)
  29. from aphrodite.connections import global_http_connection
  30. from aphrodite.distributed import (destroy_distributed_environment,
  31. destroy_model_parallel)
  32. from aphrodite.inputs import (ExplicitEncoderDecoderPrompt, TextPrompt,
  33. to_enc_dec_tuple_list, zip_enc_dec_prompts)
  34. _TEST_DIR = os.path.dirname(__file__)
  35. _TEST_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "example.txt")]
  36. _LONG_PROMPTS = [os.path.join(_TEST_DIR, "prompts", "summary.txt")]
  37. def _read_prompts(filename: str) -> List[str]:
  38. with open(filename, "r") as f:
  39. prompts = f.readlines()
  40. return prompts
  41. class _ImageAssetPrompts(TypedDict):
  42. stop_sign: str
  43. cherry_blossom: str
  44. if sys.version_info < (3, 9):
  45. # UserList cannot be subscripted
  46. class _ImageAssetsBase(UserList):
  47. pass
  48. else:
  49. class _ImageAssetsBase(UserList[ImageAsset]):
  50. pass
  51. class _ImageAssets(_ImageAssetsBase):
  52. def __init__(self) -> None:
  53. super().__init__([
  54. ImageAsset("stop_sign"),
  55. ImageAsset("cherry_blossom"),
  56. ])
  57. def prompts(self, prompts: _ImageAssetPrompts) -> List[str]:
  58. """
  59. Convenience method to define the prompt for each test image.
  60. The order of the returned prompts matches the order of the
  61. assets when iterating through this object.
  62. """
  63. return [prompts["stop_sign"], prompts["cherry_blossom"]]
  64. IMAGE_ASSETS = _ImageAssets()
  65. """Singleton instance of :class:`_ImageAssets`."""
  66. @pytest.fixture(autouse=True)
  67. def init_test_http_connection():
  68. # pytest_asyncio may use a different event loop per test
  69. # so we need to make sure the async client is created anew
  70. global_http_connection.reuse_client = False
  71. def cleanup():
  72. destroy_model_parallel()
  73. destroy_distributed_environment()
  74. with contextlib.suppress(AssertionError):
  75. torch.distributed.destroy_process_group()
  76. gc.collect()
  77. if not is_cpu():
  78. torch.cuda.empty_cache()
  79. @pytest.fixture()
  80. def should_do_global_cleanup_after_test(request) -> bool:
  81. """Allow subdirectories to skip global cleanup by overriding this fixture.
  82. This can provide a ~10x speedup for non-GPU unit tests since they don't need
  83. to initialize torch.
  84. """
  85. if request.node.get_closest_marker("skip_global_cleanup"):
  86. return False
  87. return True
  88. @pytest.fixture(autouse=True)
  89. def cleanup_fixture(should_do_global_cleanup_after_test: bool):
  90. yield
  91. if should_do_global_cleanup_after_test:
  92. cleanup()
  93. @pytest.fixture
  94. def example_prompts() -> List[str]:
  95. prompts = []
  96. for filename in _TEST_PROMPTS:
  97. prompts += _read_prompts(filename)
  98. return prompts
  99. class DecoderPromptType(Enum):
  100. """For encoder/decoder models only."""
  101. CUSTOM = 1
  102. NONE = 2
  103. EMPTY_STR = 3
  104. @pytest.fixture
  105. def example_encoder_decoder_prompts(
  106. ) -> Dict[DecoderPromptType, List[ExplicitEncoderDecoderPrompt]]:
  107. '''
  108. Returns an encoder prompt list and a decoder prompt list, wherein each pair
  109. of same-index entries in both lists corresponds to an (encoder prompt,
  110. decoder prompt) tuple.
  111. Returns:
  112. * Encoder prompt list
  113. * Decoder prompt list (reverse of encoder prompt list)
  114. '''
  115. encoder_prompts = []
  116. for filename in _TEST_PROMPTS:
  117. encoder_prompts += _read_prompts(filename)
  118. custom_decoder_prompts = encoder_prompts[::-1]
  119. empty_str_decoder_prompts = [""] * len(encoder_prompts)
  120. none_decoder_prompts = [None] * len(encoder_prompts)
  121. # NONE decoder prompt type
  122. return {
  123. DecoderPromptType.NONE:
  124. zip_enc_dec_prompts(encoder_prompts, none_decoder_prompts),
  125. DecoderPromptType.EMPTY_STR:
  126. zip_enc_dec_prompts(encoder_prompts, empty_str_decoder_prompts),
  127. DecoderPromptType.CUSTOM:
  128. zip_enc_dec_prompts(encoder_prompts, custom_decoder_prompts),
  129. }
  130. @pytest.fixture
  131. def example_long_prompts() -> List[str]:
  132. prompts = []
  133. for filename in _LONG_PROMPTS:
  134. prompts += _read_prompts(filename)
  135. return prompts
  136. @pytest.fixture(scope="session")
  137. def image_assets() -> _ImageAssets:
  138. return IMAGE_ASSETS
  139. _T = TypeVar("_T", nn.Module, torch.Tensor, BatchEncoding, BatchFeature)
  140. class HfRunner:
  141. def wrap_device(self, input: _T) -> _T:
  142. if not is_cpu():
  143. return input.to("cuda")
  144. else:
  145. return input.to("cpu")
  146. def __init__(
  147. self,
  148. model_name: str,
  149. dtype: str = "half",
  150. *,
  151. model_kwargs: Optional[Dict[str, Any]] = None,
  152. is_embedding_model: bool = False,
  153. auto_cls=AutoModelForCausalLM,
  154. postprocess_inputs: Callable[[BatchEncoding],
  155. BatchEncoding] = identity,
  156. ) -> None:
  157. torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
  158. self.model_name = model_name
  159. if is_embedding_model:
  160. # Lazy init required for AMD CI
  161. from sentence_transformers import SentenceTransformer
  162. self.model = self.wrap_device(
  163. SentenceTransformer(
  164. model_name,
  165. device="cpu",
  166. ).to(dtype=torch_dtype))
  167. else:
  168. model_kwargs = model_kwargs if model_kwargs is not None else {}
  169. self.model = self.wrap_device(
  170. auto_cls.from_pretrained(
  171. model_name,
  172. torch_dtype=torch_dtype,
  173. trust_remote_code=True,
  174. **model_kwargs,
  175. ))
  176. self.tokenizer = AutoTokenizer.from_pretrained(
  177. model_name,
  178. torch_dtype=torch_dtype,
  179. trust_remote_code=True,
  180. )
  181. try:
  182. # don't put this import at the top level
  183. # it will call torch.cuda.device_count()
  184. from transformers import AutoProcessor # noqa: F401
  185. self.processor = AutoProcessor.from_pretrained(
  186. model_name,
  187. torch_dtype=torch_dtype,
  188. trust_remote_code=True,
  189. )
  190. except Exception as exc:
  191. logger.warning(
  192. "Unable to auto-load HuggingFace processor for model (%s). "
  193. "Using tokenizer instead. Reason: %s", model_name, exc)
  194. self.processor = self.tokenizer
  195. self.postprocess_inputs = postprocess_inputs
  196. def generate(
  197. self,
  198. prompts: List[str],
  199. images: Optional[List[Image.Image]] = None,
  200. **kwargs: Any,
  201. ) -> List[Tuple[List[List[int]], List[str]]]:
  202. if images:
  203. assert len(prompts) == len(images)
  204. outputs: List[Tuple[List[List[int]], List[str]]] = []
  205. for i, prompt in enumerate(prompts):
  206. processor_kwargs: Dict[str, Any] = {
  207. "text": prompt,
  208. "return_tensors": "pt",
  209. }
  210. if images is not None and images[i] is not None:
  211. processor_kwargs["images"] = images[i]
  212. inputs = self.processor(**processor_kwargs)
  213. inputs = self.postprocess_inputs(inputs)
  214. output_ids = self.model.generate(
  215. **self.wrap_device(inputs),
  216. use_cache=True,
  217. **kwargs,
  218. )
  219. output_str = self.processor.batch_decode(
  220. output_ids,
  221. skip_special_tokens=True,
  222. clean_up_tokenization_spaces=False,
  223. )
  224. output_ids = output_ids.cpu().tolist()
  225. outputs.append((output_ids, output_str))
  226. return outputs
  227. def generate_greedy(
  228. self,
  229. prompts: List[str],
  230. max_tokens: int,
  231. images: Optional[List[Image.Image]] = None,
  232. **kwargs: Any,
  233. ) -> List[Tuple[List[int], str]]:
  234. outputs = self.generate(prompts,
  235. do_sample=False,
  236. max_new_tokens=max_tokens,
  237. images=images,
  238. **kwargs)
  239. return [(output_ids[0], output_str[0])
  240. for output_ids, output_str in outputs]
  241. def generate_beam_search(
  242. self,
  243. prompts: List[str],
  244. beam_width: int,
  245. max_tokens: int,
  246. ) -> List[Tuple[List[List[int]], List[str]]]:
  247. outputs = self.generate(prompts,
  248. do_sample=False,
  249. max_new_tokens=max_tokens,
  250. num_beams=beam_width,
  251. num_return_sequences=beam_width)
  252. for i in range(len(outputs)):
  253. output_ids, output_str = outputs[i]
  254. for j in range(len(output_ids)):
  255. output_ids[j] = [
  256. x for x in output_ids[j]
  257. if x != self.tokenizer.pad_token_id
  258. ]
  259. outputs[i] = (output_ids, output_str)
  260. return outputs
  261. def generate_greedy_logprobs(
  262. self,
  263. prompts: List[str],
  264. max_tokens: int,
  265. images: Optional[List[Image.Image]] = None,
  266. **kwargs: Any,
  267. ) -> List[List[torch.Tensor]]:
  268. all_logprobs: List[List[torch.Tensor]] = []
  269. for i, prompt in enumerate(prompts):
  270. processor_kwargs: Dict[str, Any] = {
  271. "text": prompt,
  272. "return_tensors": "pt",
  273. }
  274. if images is not None and images[i] is not None:
  275. processor_kwargs["images"] = images[i]
  276. inputs = self.processor(**processor_kwargs)
  277. inputs = self.postprocess_inputs(inputs)
  278. output = self.model.generate(
  279. **self.wrap_device(inputs),
  280. use_cache=True,
  281. do_sample=False,
  282. max_new_tokens=max_tokens,
  283. output_hidden_states=True,
  284. return_dict_in_generate=True,
  285. **kwargs,
  286. )
  287. seq_logprobs: List[torch.Tensor] = []
  288. for hidden_states in output.hidden_states:
  289. last_hidden_states = hidden_states[-1][0]
  290. logits = torch.matmul(
  291. last_hidden_states,
  292. self.model.get_output_embeddings().weight.t(),
  293. )
  294. if self.model.get_output_embeddings().bias is not None:
  295. logits += self.model.get_output_embeddings(
  296. ).bias.unsqueeze(0)
  297. logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
  298. seq_logprobs.append(logprobs)
  299. all_logprobs.append(seq_logprobs)
  300. return all_logprobs
  301. def _hidden_states_to_logprobs(
  302. self,
  303. hidden_states,
  304. num_logprobs,
  305. ) -> Tuple[List[Dict[int, float]], int]:
  306. seq_logprobs: List[torch.Tensor] = []
  307. output_len = len(hidden_states)
  308. for _, hidden_state in enumerate(hidden_states):
  309. last_hidden_states = hidden_state[-1][0]
  310. logits = torch.matmul(
  311. last_hidden_states,
  312. self.model.get_output_embeddings().weight.t(),
  313. )
  314. if getattr(self.model.get_output_embeddings(), "bias",
  315. None) is not None:
  316. logits += self.model.get_output_embeddings().bias.unsqueeze(0)
  317. logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
  318. seq_logprobs.append(logprobs)
  319. # convert to dict
  320. seq_logprobs_lst: List[Dict[int, float]] = []
  321. for tok_idx, tok_logprobs in enumerate(seq_logprobs):
  322. # drop prompt logprobs
  323. if tok_idx == 0:
  324. tok_logprobs = tok_logprobs[-1, :].reshape(1, -1)
  325. topk = tok_logprobs.topk(num_logprobs)
  326. tok_logprobs_dct = {}
  327. for token_id, logprob in zip(topk.indices[0], topk.values[0]):
  328. tok_logprobs_dct[token_id.item()] = logprob.item()
  329. seq_logprobs_lst.append(tok_logprobs_dct)
  330. return (
  331. seq_logprobs_lst,
  332. output_len,
  333. )
  334. def generate_greedy_logprobs_limit(
  335. self,
  336. prompts: List[str],
  337. max_tokens: int,
  338. num_logprobs: int,
  339. images: Optional[List[Image.Image]] = None,
  340. audios: Optional[List[Tuple[np.ndarray, int]]] = None,
  341. **kwargs: Any,
  342. ) -> List[Tuple[List[int], str, List[Dict[int, float]]]]:
  343. all_logprobs: List[List[Dict[int, float]]] = []
  344. all_output_ids: List[List[int]] = []
  345. all_output_strs: List[str] = []
  346. for i, prompt in enumerate(prompts):
  347. processor_kwargs: Dict[str, Any] = {
  348. "text": prompt,
  349. "return_tensors": "pt",
  350. }
  351. if images is not None and images[i] is not None:
  352. processor_kwargs["images"] = images[i]
  353. if audios is not None:
  354. audio, sr = audios[i]
  355. processor_kwargs["audio"] = audio
  356. processor_kwargs["sampling_rate"] = sr
  357. inputs = self.processor(**processor_kwargs)
  358. inputs = self.postprocess_inputs(inputs)
  359. output = self.model.generate(
  360. **self.wrap_device(inputs),
  361. use_cache=True,
  362. do_sample=False,
  363. max_new_tokens=max_tokens,
  364. output_hidden_states=True,
  365. return_dict_in_generate=True,
  366. **kwargs,
  367. )
  368. (
  369. seq_logprobs_lst,
  370. output_len,
  371. ) = self._hidden_states_to_logprobs(output.hidden_states,
  372. num_logprobs)
  373. all_logprobs.append(seq_logprobs_lst)
  374. seq_ids = output.sequences[0]
  375. output_len = len(seq_logprobs_lst)
  376. output_ids = seq_ids[-output_len:]
  377. all_output_ids.append(output_ids.tolist())
  378. all_output_strs.append(self.tokenizer.decode(output_ids))
  379. outputs = zip(all_output_ids, all_output_strs, all_logprobs)
  380. return [(output_ids, output_str, output_logprobs)
  381. for output_ids, output_str, output_logprobs in outputs]
  382. def generate_encoder_decoder_greedy_logprobs_limit(
  383. self,
  384. encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
  385. max_tokens: int,
  386. num_logprobs: int,
  387. **kwargs: Any,
  388. ) -> List[Tuple[List[int], str, List[Dict[int, float]]]]:
  389. '''
  390. Greedy logprobs generation for Aphrodite encoder/decoder models
  391. '''
  392. all_logprobs: List[List[Dict[int, float]]] = []
  393. all_output_ids: List[List[int]] = []
  394. all_output_strs: List[str] = []
  395. for (encoder_prompt,
  396. decoder_prompt) in to_enc_dec_tuple_list(encoder_decoder_prompts):
  397. encoder_input_ids = self.wrap_device(
  398. self.tokenizer(encoder_prompt, return_tensors="pt").input_ids)
  399. decoder_input_ids = (
  400. None if decoder_prompt is None else self.wrap_device(
  401. self.tokenizer(decoder_prompt,
  402. return_tensors="pt").input_ids))
  403. output = self.model.generate(
  404. encoder_input_ids,
  405. decoder_input_ids=decoder_input_ids,
  406. use_cache=True,
  407. do_sample=False,
  408. max_new_tokens=max_tokens,
  409. output_hidden_states=True,
  410. return_dict_in_generate=True,
  411. **kwargs,
  412. )
  413. (
  414. seq_logprobs_lst,
  415. output_len,
  416. ) = self._hidden_states_to_logprobs(output.decoder_hidden_states,
  417. num_logprobs)
  418. all_logprobs.append(seq_logprobs_lst)
  419. seq_ids = output.sequences[0]
  420. output_ids = seq_ids[-output_len:]
  421. all_output_ids.append(output_ids.tolist())
  422. all_output_strs.append(self.tokenizer.decode(output_ids))
  423. outputs = zip(all_output_ids, all_output_strs, all_logprobs)
  424. return [(output_ids, output_str, output_logprobs)
  425. for output_ids, output_str, output_logprobs in outputs]
  426. def encode(self, prompts: List[str]) -> List[List[torch.Tensor]]:
  427. return self.model.encode(prompts)
  428. def __enter__(self):
  429. return self
  430. def __exit__(self, exc_type, exc_value, traceback):
  431. del self.model
  432. cleanup()
  433. @pytest.fixture(scope="session")
  434. def hf_runner():
  435. return HfRunner
  436. class AphroditeRunner:
  437. def __init__(
  438. self,
  439. model_name: str,
  440. tokenizer_name: Optional[str] = None,
  441. # Use smaller max model length, otherwise bigger model cannot run due
  442. # to kv cache size limit.
  443. max_model_len: int = 1024,
  444. dtype: str = "half",
  445. disable_log_stats: bool = True,
  446. tensor_parallel_size: int = 1,
  447. block_size: int = 16,
  448. enable_chunked_prefill: bool = False,
  449. swap_space: int = 4,
  450. enforce_eager: Optional[bool] = False,
  451. **kwargs,
  452. ) -> None:
  453. self.model = LLM(
  454. model=model_name,
  455. tokenizer=tokenizer_name,
  456. trust_remote_code=True,
  457. dtype=dtype,
  458. swap_space=swap_space,
  459. enforce_eager=enforce_eager,
  460. disable_log_stats=disable_log_stats,
  461. tensor_parallel_size=tensor_parallel_size,
  462. max_model_len=max_model_len,
  463. block_size=block_size,
  464. enable_chunked_prefill=enable_chunked_prefill,
  465. **kwargs,
  466. )
  467. def generate(
  468. self,
  469. prompts: List[str],
  470. sampling_params: SamplingParams,
  471. images: Optional[Union[List[Image.Image],
  472. List[List[Image.Image]]]] = None,
  473. ) -> List[Tuple[List[List[int]], List[str]]]:
  474. if images is not None:
  475. assert len(prompts) == len(images)
  476. inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
  477. if images is not None:
  478. for i, image in enumerate(images):
  479. inputs[i]["multi_modal_data"] = {"image": image}
  480. req_outputs = self.model.generate(inputs,
  481. sampling_params=sampling_params)
  482. outputs: List[Tuple[List[List[int]], List[str]]] = []
  483. for req_output in req_outputs:
  484. prompt_str = req_output.prompt
  485. prompt_ids = req_output.prompt_token_ids
  486. req_sample_output_ids: List[List[int]] = []
  487. req_sample_output_strs: List[str] = []
  488. for sample in req_output.outputs:
  489. output_str = sample.text
  490. output_ids = list(sample.token_ids)
  491. req_sample_output_ids.append(prompt_ids + output_ids)
  492. req_sample_output_strs.append(prompt_str + output_str)
  493. outputs.append((req_sample_output_ids, req_sample_output_strs))
  494. return outputs
  495. def _final_steps_generate_w_logprobs(
  496. self,
  497. req_outputs: List[RequestOutput],
  498. ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
  499. outputs: List[Tuple[List[int], str, Optional[SampleLogprobs]]] = []
  500. for req_output in req_outputs:
  501. for sample in req_output.outputs:
  502. output_str = sample.text
  503. output_ids = list(sample.token_ids)
  504. output_logprobs = sample.logprobs
  505. outputs.append((output_ids, output_str, output_logprobs))
  506. return outputs
  507. def generate_w_logprobs(
  508. self,
  509. prompts: List[str],
  510. sampling_params: SamplingParams,
  511. images: Optional[Union[List[Image.Image],
  512. List[List[Image.Image]]]] = None,
  513. audios: Optional[Union[List[Tuple[np.ndarray, int]],
  514. List[List[Tuple[np.ndarray, int]]]]] = None
  515. ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
  516. assert sampling_params.logprobs is not None
  517. if images is not None:
  518. assert len(prompts) == len(images)
  519. inputs = [TextPrompt(prompt=prompt) for prompt in prompts]
  520. if images is not None:
  521. for i, image in enumerate(images):
  522. inputs[i]["multi_modal_data"] = {"image": image}
  523. if audios is not None:
  524. for i, audio in enumerate(audios):
  525. inputs[i]["multi_modal_data"] = {"audio": audio}
  526. req_outputs = self.model.generate(inputs,
  527. sampling_params=sampling_params)
  528. return self._final_steps_generate_w_logprobs(req_outputs)
  529. def generate_encoder_decoder_w_logprobs(
  530. self,
  531. encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
  532. sampling_params: SamplingParams,
  533. ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
  534. '''
  535. Logprobs generation for Aphrodite encoder/decoder models
  536. '''
  537. assert sampling_params.logprobs is not None
  538. req_outputs = self.model.generate(encoder_decoder_prompts,
  539. sampling_params=sampling_params)
  540. return self._final_steps_generate_w_logprobs(req_outputs)
  541. def generate_greedy(
  542. self,
  543. prompts: List[str],
  544. max_tokens: int,
  545. images: Optional[List[Image.Image]] = None,
  546. ) -> List[Tuple[List[int], str]]:
  547. greedy_params = SamplingParams(temperature=0.0, max_tokens=max_tokens)
  548. outputs = self.generate(prompts, greedy_params, images=images)
  549. return [(output_ids[0], output_str[0])
  550. for output_ids, output_str in outputs]
  551. def generate_greedy_logprobs(
  552. self,
  553. prompts: List[str],
  554. max_tokens: int,
  555. num_logprobs: int,
  556. images: Optional[Union[List[Image.Image],
  557. List[List[Image.Image]]]] = None,
  558. audios: Optional[Union[List[Tuple[np.ndarray, int]],
  559. List[List[Tuple[np.ndarray, int]]]]] = None,
  560. stop_token_ids: Optional[List[int]] = None,
  561. ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
  562. greedy_logprobs_params = SamplingParams(temperature=0.0,
  563. max_tokens=max_tokens,
  564. logprobs=num_logprobs,
  565. stop_token_ids=stop_token_ids)
  566. outputs = self.generate_w_logprobs(prompts,
  567. greedy_logprobs_params,
  568. images=images,
  569. audios=audios)
  570. return [(output_ids, output_str, output_logprobs)
  571. for output_ids, output_str, output_logprobs in outputs]
  572. def generate_encoder_decoder_greedy_logprobs(
  573. self,
  574. encoder_decoder_prompts: List[ExplicitEncoderDecoderPrompt[str, str]],
  575. max_tokens: int,
  576. num_logprobs: int,
  577. ) -> List[Tuple[List[int], str, Optional[SampleLogprobs]]]:
  578. greedy_logprobs_params = SamplingParams(temperature=0.0,
  579. use_beam_search=False,
  580. max_tokens=max_tokens,
  581. logprobs=num_logprobs)
  582. '''
  583. Greedy logprobs generation for Aphrodite encoder/decoder models
  584. '''
  585. outputs = self.generate_encoder_decoder_w_logprobs(
  586. encoder_decoder_prompts, greedy_logprobs_params)
  587. return [(output_ids, output_str, output_logprobs)
  588. for output_ids, output_str, output_logprobs in outputs]
  589. def generate_beam_search(
  590. self,
  591. prompts: List[str],
  592. beam_width: int,
  593. max_tokens: int,
  594. ) -> List[Tuple[List[List[int]], List[str]]]:
  595. beam_search_params = SamplingParams(n=beam_width,
  596. use_beam_search=True,
  597. temperature=0.0,
  598. max_tokens=max_tokens)
  599. outputs = self.generate(prompts, beam_search_params)
  600. return outputs
  601. def encode(self, prompts: List[str]) -> List[List[float]]:
  602. req_outputs = self.model.encode(prompts)
  603. outputs = []
  604. for req_output in req_outputs:
  605. embedding = req_output.outputs.embedding
  606. outputs.append(embedding)
  607. return outputs
  608. def __enter__(self):
  609. return self
  610. def __exit__(self, exc_type, exc_value, traceback):
  611. del self.model
  612. cleanup()
  613. @pytest.fixture(scope="session")
  614. def aphrodite_runner():
  615. return AphroditeRunner
  616. def get_tokenizer_pool_config(tokenizer_group_type):
  617. if tokenizer_group_type is None:
  618. return None
  619. if tokenizer_group_type == "ray":
  620. return TokenizerPoolConfig(pool_size=1,
  621. pool_type="ray",
  622. extra_config={})
  623. if isinstance(tokenizer_group_type, type):
  624. return TokenizerPoolConfig(pool_size=1,
  625. pool_type=tokenizer_group_type,
  626. extra_config={})
  627. raise ValueError(f"Unknown tokenizer_group_type: {tokenizer_group_type}")
  628. @pytest.fixture()
  629. def temporary_enable_log_propagate():
  630. import logging
  631. logger = logging.getLogger("aphrodite")
  632. logger.propagate = True
  633. yield
  634. logger.propagate = False
  635. @pytest.fixture()
  636. def caplog_aphrodite(temporary_enable_log_propagate, caplog):
  637. # To capture aphrodite log, we should enable propagate=True temporarily
  638. # because caplog depends on logs propagated to the root logger.
  639. yield caplog
  640. @pytest.fixture(scope="session")
  641. def num_gpus_available():
  642. """Get number of GPUs without initializing the CUDA context
  643. in current process."""
  644. return cuda_device_count_stateless()
  645. temp_dir = tempfile.gettempdir()
  646. _dummy_path = os.path.join(temp_dir, "dummy_opt")
  647. @pytest.fixture
  648. def dummy_opt_path():
  649. json_path = os.path.join(_dummy_path, "config.json")
  650. if not os.path.exists(_dummy_path):
  651. snapshot_download(repo_id="facebook/opt-125m",
  652. local_dir=_dummy_path,
  653. ignore_patterns=[
  654. "*.bin", "*.bin.index.json", "*.pt", "*.h5",
  655. "*.msgpack"
  656. ])
  657. assert os.path.exists(json_path)
  658. with open(json_path, "r") as f:
  659. config = json.load(f)
  660. config["architectures"] = ["MyOPTForCausalLM"]
  661. with open(json_path, "w") as f:
  662. json.dump(config, f)
  663. return _dummy_path