Skip to content

Instantly share code, notes, and snippets.

@mrmaheshrajput
Created March 26, 2024 14:28
Show Gist options
  • Save mrmaheshrajput/6054f5b3a2830490556256e78fece925 to your computer and use it in GitHub Desktop.
Save mrmaheshrajput/6054f5b3a2830490556256e78fece925 to your computer and use it in GitHub Desktop.
import evaluate
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from datasets import load_dataset
# Load the fine-tuned model and tokenizer
model_name = "your-fine-tuned-model-name"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
# Load the test dataset
test_dataset = load_dataset("your-test-dataset-name")
# Function to generate predictions from model
def generate_predictions(batch):
inputs = tokenizer(batch["input"], padding="max_length", truncation=True, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(**inputs)
return tokenizer.batch_decode(outputs, skip_special_tokens=True)
# Generate predictions on the test dataset
predictions = test_dataset.map(generate_predictions, batched=True, batch_size=4)
# Convert predictions and references to lists
predictions_list = predictions["predictions"]
references_list = predictions["references"]
# Compute ROUGE scores
rouge_scores = evaluate.compute(predictions_list, references_list)
# Print ROUGE scores
print(rouge_scores)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment