import torch
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
model_id = "Arthur-75/quester-qwen3-4B"
system_prompt = "Generate relevant single-word keywords to improve retrieval performance. Only output unique keywords, separated by commas."
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_id,
trust_remote_code=True,
)
model.eval()
def clean_output(text: str) -> str:
text = text.strip()
if "</think>" in text:
text = text.split("</think>", 1)[-1].strip()
if "<think>" in text:
text = text.split("<think>", 1)[-1].strip()
if text.lower().startswith("keywords:"):
text = text[len("keywords:"):].strip()
if text.lower().startswith("keyword:"):
text = text[len("keyword:"):].strip()
if "\nassistant" in text:
text = text.split("\nassistant", 1)[-1].strip()
if "assistant\n" in text:
text = text.split("assistant\n", 1)[-1].strip()
return text
def generate_keywords(query: str) -> str:
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content":"[QUERY]: "+ query+" [KEYWORDS]:"},
]
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=False,
)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.inference_mode():
output = model.generate(
**inputs,
max_new_tokens=64,
do_sample=False,
)
new_tokens = output[:, inputs["input_ids"].shape[1]:]
text = tokenizer.decode(new_tokens[0], skip_special_tokens=True)
return clean_output(text)
query = "What are the symptoms of vitamin D deficiency?"
print(generate_keywords(query))
