File size: 369 Bytes
97531f4
 
 
 
 
 
736e7a3
97531f4
9eca55e
 
97531f4
1
2
3
4
5
6
7
8
9
10
11
12
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "google/gemma-2b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)


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