mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
import os
|
|
|
|
import torch
|
|
from vllm import LLM, SamplingParams
|
|
|
|
os.environ["VLLM_USE_V1"] = "0" # vLLM V1 does not support logits processors.
|
|
|
|
|
|
class LogitsSpy:
|
|
def __init__(self):
|
|
self.processed_logits: list[torch.Tensor] = []
|
|
|
|
def __call__(self, token_ids: list[int], logits: torch.Tensor):
|
|
self.processed_logits.append(logits)
|
|
return logits
|
|
|
|
|
|
if __name__ == "__main__":
|
|
llm = LLM(model="google/gemma-2-2b-it", trust_remote_code=True)
|
|
prompts = [
|
|
"What is the capital of France?",
|
|
"Explain the theory of relativity in simple terms.",
|
|
"Write a short story about a robot learning to love.",
|
|
]
|
|
logits_spy = LogitsSpy()
|
|
sampling_params = SamplingParams(
|
|
temperature=0.7, top_p=0.95, max_tokens=5, logits_processors=[logits_spy]
|
|
)
|
|
outputs = llm.generate(prompts, sampling_params)
|
|
logits = logits_spy.processed_logits
|
|
print(f"Logits: {logits}")
|
|
print(len(logits))
|
|
print(logits[0].shape)
|
|
print(logits[1].shape)
|
|
print(logits[2].shape)
|
|
print(logits[3].shape)
|
|
print(logits[4].shape)
|
|
|
|
breakpoint()
|