src/brain.py

from transformers import GPT2LMHeadModel, GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
model = GPT2LMHeadModel.from_pretrained("distilgpt2")

def generate_answers(query):
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    max_length = input_ids.shape[1] + 100
    generated_ids = model.generate(input_ids, max_length=max_length)
    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return generated_text