File size: 429 Bytes
78047da
 
1eec729
c65051f
 
78047da
 
 
71f21fd
 
 
1
2
3
4
5
6
7
8
9
10
11
12
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "meta-llama/Llama-2-7b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)


def generate_answers(query):
    input_ids = tokenizer(query, return_tensors="pt").input_ids
    output = model.generate(input_ids, max_new_tokens=100)
    return tokenizer.decode(output[0], skip_special_tokens=True)