generator.py

from transformers import AutoTokenizer
from transformers import AutoModelForQuestionAnswering
from transformers import QuestionAnsweringPipeline

model_name = 'KoichiYasuoka/bert-base-japanese-wikipedia-ud-head'

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
qap = QuestionAnsweringPipeline(tokenizer=tokenizer, model=model)

def generate_response(st, prompt):
    # Transformersで回答を作成
    answer = qap(context=st.session_state.content[:100], question=prompt)

    return answer["answer"]