from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
import torch
import torch_npu
import argparse
import os
from openmind_hub import snapshot_download

# 设置环境变量
os.environ['DEFAULT_DOWNLOAD_TIMEOUT'] = "600"
os.environ['DEFAULT_REQUEST_TIMEOUT'] = "600"

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Jinan_AICC/gelectra-base-germanquad-distilled",
        default=None,
    )
    args = parser.parse_args()
    return args

args = parse_args()

if args.model_name_or_path:
    model_path = args.model_name_or_path
else:
    model_path = snapshot_download(
        "Jinan_AICC/gelectra-base-germanquad-distilled",
        revision="main",
        ignore_patterns=["*.h5", "*.ot", "*.msgpack"],       
    )

# a) Get predictions
nlp = pipeline('question-answering', model=model_path, tokenizer=model_path)
QA_input = {
    'question': 'Why is model conversion important?',
    'context': 'The option to convert models between FARM and transformers gives freedom to the user and let people easily switch between frameworks.'
}
res = nlp(QA_input)

# b) Load model & tokenizer
model = AutoModelForQuestionAnswering.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)