Search code examples
pythonnlpartificial-intelligencetorchbert-language-model

Fine-Tuned ALBERT Question and Answering with HuggingFace


I'm trying to create a question and answering AI, I would like it to be as accurate as possible without having to train the model myself.

I can create a simple AI using the existing base models like so via their documentation:

from transformers import AlbertTokenizer, AlbertForQuestionAnswering
import torch
tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
model = AlbertForQuestionAnswering.from_pretrained('albert-base-v2')
question, text = "What does He like?", "He likes bears"
inputs = tokenizer(question, text, return_tensors='pt')
start_positions = torch.tensor([1])
end_positions = torch.tensor([3])
outputs = model(**inputs, start_positions=start_positions, end_positions=end_positions)
loss = outputs.loss
start_scores = outputs.start_logits
end_scores = outputs.end_logits

answer_start = torch.argmax(start_scores)  # get the most likely beginning of answer with the argmax of the score
answer_end = torch.argmax(end_scores) + 1
tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))

However this model doesn't answer questions as accurate as others. On the HuggingFace site I've found an example that I'd like to use of a fine-tuned model

However the instructions show how to train a model like so. The example works on the page so clearly a pretrained model of the exists.

Does anyone know how I can reuse the existing models so I don't have to train one from scratch?


Solution

  • Turns out I just needed to grab an additional identifier when trying to request the model:

    from transformers import AlbertTokenizer, AlbertForQuestionAnswering
    import torch
    
    MODEL_PATH = 'ktrapeznikov/albert-xlarge-v2-squad-v2';
    
    tokenizer = AlbertTokenizer.from_pretrained(MODEL_PATH)
    model = AlbertForQuestionAnswering.from_pretrained(MODEL_PATH)
    

    For future reference this information can be grabbed from the transformers use button. Seem in the image below.

    enter image description here