File size: 562 Bytes
1262e09
b985953
89bc9d2
 
622f41b
b985953
89bc9d2
1262e09
 
 
1
2
3
4
5
6
7
8
9
10
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

def generate_answer(context, question):
    prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
    inputs = tokenizer.encode(prompt, return_tensors='pt', truncation=True, max_length=1024)
    outputs = model.generate(inputs, max_new_tokens=60, do_sample=True, pad_token_id=tokenizer.eos_token_id)
    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()