conftest.py 27 KB

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