The model requires specific prompt construction to trigger its 'Knowledge Retrieval' mode:
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4model_id = 'solvrays/scribegene-llm-v1.1'
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForCausalLM.from_pretrained(
7 model_id,
8 device_map='auto',
9 torch_dtype=torch.bfloat16,
10 quantization_config={'load_in_4bit': True}
11)
12
13def query_model(user_query):
14 # High-Precision Retrieval Template
15 prompt = f'### Knowledge Retrieval Content: {user_query}\n### Verified Response: '
16 inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
17 outputs = model.generate(**inputs, max_new_tokens=512, do_sample=False)
18 return tokenizer.decode(outputs[0], skip_special_tokens=True).split('### Verified Response:')[-1].strip()