!pip install --upgrade --quiet "google-cloud-aiplatform[evaluation]"  huggingface_hub transformers datasets

%env PROJECT_ID=gcp-partnership-412108
%env LOCATION=us-central1
%env CONTAINER_URI=us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-2.ubuntu2204.py310

!gcloud auth login
!gcloud auth application-default login  # For local development
!gcloud config set project $PROJECT_ID

!gcloud services enable aiplatform.googleapis.com
!gcloud services enable compute.googleapis.com
!gcloud services enable container.googleapis.com
!gcloud services enable containerregistry.googleapis.com
!gcloud services enable containerfilesystem.googleapis.com

from huggingface_hub import interpreter_login

interpreter_login()

import os
from google.cloud import aiplatform

aiplatform.init(
    project=os.getenv("PROJECT_ID"),
    location=os.getenv("LOCATION"),
)

from huggingface_hub import get_token

vertex_model_name = "llama-3-1-8b-instruct"

model = aiplatform.Model.upload(
    display_name=vertex_model_name,
    serving_container_image_uri=os.getenv("CONTAINER_URI"),
    serving_container_environment_variables={
        "MODEL_ID": "meta-llama/Meta-Llama-3.1-8B-Instruct",
        "MAX_INPUT_TOKENS": "8000",
        "MAX_TOTAL_TOKENS": "8192",
        "MAX_BATCH_PREFILL_TOKENS": "8192",
        "HUGGING_FACE_HUB_TOKEN": get_token(),
    },
    serving_container_ports=[8080],
)
model.wait()  # wait for the model to be registered

# create endpoint
endpoint = aiplatform.Endpoint.create(display_name=f"{vertex_model_name}-endpoint")

# deploy model to 1x NVIDIA L4
deployed_model = model.deploy(
    endpoint=endpoint,
    machine_type="g2-standard-4",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
)

import re
from transformers import AutoTokenizer

# grep the model id from the container spec environment variables
model_id = next(
    (
        re.search(r'value: "(.+)"', str(item)).group(1)
        for item in list(model.container_spec.env)
        if "MODEL_ID" in str(item)
    ),
    None,
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")

generation_config = {
    "max_new_tokens": 256,
    "do_sample": True,
    "top_p": 0.2,
    "temperature": 0.2,
}


def generate(prompt, generation_config=generation_config):
    formatted_prompt = tokenizer.apply_chat_template(
        [
            {"role": "user", "content": prompt},
        ],
        tokenize=False,
        add_generation_prompt=True,
    )

    payload = {"inputs": formatted_prompt, "parameters": generation_config}
    output = deployed_model.predict(instances=[payload])
    generated_text = output.predictions[0]
    return generated_text


generate("How many people live in Berlin?", generation_config)
# 'The population of Berlin is approximately 6.578 million as of my cut off data. However, considering it provides real-time updates, the current population might be slightly higher'

from vertexai.evaluation import EvalTask, PointwiseMetric

g_eval_coherence = """
You are an expert evaluator. You will be given one summary written for a news article.
Your task is to rate the summary on one metric.
Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed.

Evaluation Criteria:

Coherence (1-5) - the collective quality of all sentences. We align this dimension with the DUC quality question of structure and coherence whereby "the summary should be well-structured and well-organized. The summary should not just be a heap of related information, but should build from sentence to a coherent body of information about a topic."

Evaluation Steps:

1. Read the news article carefully and identify the main topic and key points.
2. Read the summary and compare it to the news article. Check if the summary covers the main topic and key points of the news article, and if it presents them in a clear and logical order.
3. Assign a score for coherence on a scale of 1 to 5, where 1 is the lowest and 5 is the highest based on the Evaluation Criteria.


Example:


Source Text:

{text}

Summary:

{response}

Evaluation Form (scores ONLY):

- Coherence:"""

metric = PointwiseMetric(
    metric="g-eval-coherence",
    metric_prompt_template=g_eval_coherence,
)

from datasets import load_dataset

subset_size = 15
dataset = load_dataset("argilla/news-summary", split=f"train").shuffle(seed=42).select(range(subset_size))

# print first 150 characters of the first article
print(dataset[0]["text"][:150])

# remove all columns except for "text"
to_remove = [col for col in dataset.features.keys() if col != "text"]
dataset = dataset.remove_columns(to_remove)
df = dataset.to_pandas()
df.head()

summarization_prompts = {
    "simple": "Summarize the following news article: {text}",
    "eli5": "Summarize the following news article in a way a 5 year old would understand: {text}",
    "detailed": """Summarize the given news article, text, including all key points and supporting details? The summary should be comprehensive and accurately reflect the main message and arguments presented in the original text, while also being concise and easy to understand. To ensure accuracy, please read the text carefully and pay attention to any nuances or complexities in the language.
  
Article:
{text}""",
}

import uuid


results = {}
for prompt_name, prompt in summarization_prompts.items():
    prompt = summarization_prompts[prompt_name]

    # 1. add new prompt column
    df["prompt"] = df["text"].apply(lambda x: prompt.format(text=x))

    # 2. create eval task
    eval_task = EvalTask(
        dataset=df,
        metrics=[metric],
        experiment="llama-3-1-8b-instruct",
    )
    # 3. run eval task
    # Note: If the last iteration takes > 1 minute you might need to retry the evaluation
    exp_results = eval_task.evaluate(
        model=generate, experiment_run_name=f"prompt-{prompt_name}-{str(uuid.uuid4())[:8]}"
    )
    print(f"{prompt_name}: {exp_results.summary_metrics['g-eval-coherence/mean']}")
    results[prompt_name] = exp_results.summary_metrics["g-eval-coherence/mean"]

for prompt_name, score in sorted(results.items(), key=lambda x: x[1], reverse=True):
    print(f"{prompt_name}: {score}")

deployed_model.undeploy_all()
deployed_model.delete()
model.delete()