from transformers import MBartForConditionalGeneration, MBart50Tokenizer

# Load the model and tokenizer
model_name = "LocalDoc/mbart_large_qa_azerbaijan"
tokenizer = MBart50Tokenizer.from_pretrained(model_name, src_lang="en_XX", tgt_lang="az_AZ")
model = MBartForConditionalGeneration.from_pretrained(model_name)

def answer_question(context, question):
    # Prepare input text
    input_text = f"context: {context} question: {question}"
    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True, padding="max_length")
    
    # Generate answer
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=128,
        num_beams=5,
        early_stopping=True
    )
    
    # Decode the answer
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

# Example usage
context = """Ey iman gətirənlər! Səbir edərək və namaz qılaraq Allahdan kömək istəyin! Şübhəsiz ki, Allah səbir edənlərlədir.
Ey iman gətirənlər, əgər yalnız Allaha ibadət edirsinizsə, sizə verdiyimiz ruzilərin təmiz olanlarından yeyin və Ona şükür edin!
Ey iman gətirənlər! Oruc tutmaq sizdən əvvəlkilərə fərz buyurulduğu kimi sizə də fərz buyuruldu ki, bəlkə, günahlardan çəkinəsiniz!
Ey iman gətirənlər, bütövlüklə İslama daxil olun! Şeytanın addımlarını izləməyin! Şübhəsiz ki, o sizin açıq-aydın düşməninizdir.
Ey iman gətirənlər! Həqiqətən, içki, qumar, bütlər və fal oxları şeytan əməli olan murdar şeylərdir. Bunlardan uzaq durun ki, bəlkə nicat tapasınız!
Ey iman gətirənlər! “Allah yolunda döyüşə çıxın!” - deyildiyi zaman sizə nə oldu ki, yerə yapışıb qaldınız. Axirətin əvəzinə dünya həyatına razı oldunuz?
Dünya malı axirət nemətlərinin yanında çox az bir şeydir.
Həqiqətən, insan çox zalım və çox nankordur."""

question = "Allah kiminlədir ?"

answer = answer_question(context, question)
print(answer)