File size: 1,019 Bytes
e73d95b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
from vllm import SamplingParams
from prefixLLM import PrefixLLM
from template import SYSTEM_PROMPT, SHORT_TRIGGER
if __name__ == "__main__":
llm = PrefixLLM(model="amandaa/AutoL2S-7b")
max_tokens, temp = 32768, 0.7
sampling_params_route = SamplingParams(max_tokens=max_tokens, temperature=temp, stop=["<specialLong>"], include_stop_str_in_output=True)
sampling_params_force_think = SamplingParams(max_tokens=max_tokens, temperature=temp)
question = "Convert the point $(0,3)$ in rectangular coordinates to polar coordinates. Enter your answer in the form $(r,\\theta),$ where $r > 0$ and $0 \\le \\theta < 2 \\pi.$"
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": question}
]
responses = llm.route_chat(messages=messages, sampling_params_route=sampling_params_route, sampling_params_force_think=sampling_params_force_think, use_tqdm=True, trigger_word=SHORT_TRIGGER)
print(SHORT_TRIGGER + responses[0].outputs[0].text) |