Spaces:
Paused
Paused
File size: 1,043 Bytes
9189e38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from transformers import LlamaForQuestionAnswering, LlamaTokenizer
import torch
# Load the pre-trained LLaMA-2 model and tokenizer
# model_name = "meta-llama/llama-2-question-answering"
model_name = "meta-llama/Meta-Llama-3-8B"
model = LlamaForQuestionAnswering.from_pretrained(model_name)
tokenizer = LlamaTokenizer.from_pretrained(model_name)
# Prepare a question and context
question = "What is the capital of France?"
context = "France, a country in Western Europe, is known for its medieval cities, alpine villages, and Mediterranean beaches. Its capital, Paris, is famed for its fashion, gastronomy, and culture."
inputs = tokenizer(question, context, return_tensors="pt")
# Perform the question-answering
with torch.no_grad():
outputs = model(**inputs)
answer_start_index = outputs.start_logits.argmax()
answer_end_index = outputs.end_logits.argmax()
answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs.input_ids[0][answer_start_index:answer_end_index+1]))
print(f"Answer: {answer}")
|