Views
No views yet
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)def make_inference(query, model):
prompt = """\
### Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
Categorize this question into one of this two categories:
RAG
COT
Input:
{Question}
### Response:
"""
batch = tokenizer(prompt.format(Question=query), return_tensors='pt').to("cuda")
with torch.cuda.amp.autocast():
output_tokens = model.generate(**batch, max_new_tokens=10)
return output_tokensquery = "{your_question_goes_here}"
output_tokens = make_inference(query, model)
response = tokenizer.decode(output_tokens[0])
print(response)1### Below is an instruction that describes a task. Write a response that appropriately completes the request.
2
3### Instruction:
4Categorize this question into one of this two categories:
5
6RAG
7COT
8
9Input:
10{Question}
11
12### Response:
13{Category}
14
15### End
16
17"""
