from transformers import MBartForConditionalGeneration, MBart50Tokenizer
import gradio as gr
import train as tr

'''
# 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(text, question):
    # Prepare input text
    input_text = f"context: {text} question: {question}"
    inputs = tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=False, padding="max_length")
    
    # Generate answer
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_length=1024,
        num_beams=5,
        early_stopping=True
    )
    
    # Decode the answer
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

demo = gr.Interface(
  fn=answer_question,
  inputs=["text", "text"],
  outputs=["text"]
)


#demo.launch()
tr.init()