81 lines
3.5 KiB
Python
81 lines
3.5 KiB
Python
import os
|
|
|
|
import pandas as pd
|
|
from langchain.chains import LLMChain
|
|
from langchain.llms import OpenAI
|
|
from langchain.prompts import PromptTemplate
|
|
|
|
import mlflow
|
|
|
|
assert "OPENAI_API_KEY" in os.environ, (
|
|
"Please set the OPENAI_API_KEY environment variable to run this example."
|
|
)
|
|
|
|
|
|
def build_and_evaluate_model_with_prompt(prompt_template):
|
|
mlflow.start_run()
|
|
mlflow.log_param("prompt_template", prompt_template)
|
|
# Create a news summarization model using prompt engineering with LangChain. Log the model
|
|
# to MLflow Tracking
|
|
llm = OpenAI(temperature=0.9)
|
|
prompt = PromptTemplate(input_variables=["article"], template=prompt_template)
|
|
chain = LLMChain(llm=llm, prompt=prompt)
|
|
logged_model = mlflow.langchain.log_model(chain, name="model")
|
|
|
|
# Evaluate the model on a small sample dataset
|
|
sample_data = pd.read_csv("summarization_example_data.csv")
|
|
mlflow.evaluate(
|
|
model=logged_model.model_uri,
|
|
model_type="text-summarization",
|
|
data=sample_data,
|
|
targets="highlights",
|
|
)
|
|
mlflow.end_run()
|
|
|
|
|
|
prompt_template_1 = (
|
|
"Write a summary of the following article that is between triple backticks: ```{article}```"
|
|
)
|
|
print(f"Building and evaluating model with prompt: '{prompt_template_1}'")
|
|
build_and_evaluate_model_with_prompt(prompt_template_1)
|
|
|
|
prompt_template_2 = (
|
|
"Write a summary of the following article that is between triple backticks. Be concise. Make"
|
|
" sure the summary includes important nouns and dates and keywords in the original text."
|
|
" Just return the summary. Do not include any text other than the summary: ```{article}```"
|
|
)
|
|
print(f"Building and evaluating model with prompt: '{prompt_template_2}'")
|
|
build_and_evaluate_model_with_prompt(prompt_template_2)
|
|
|
|
# Load the evaluation results
|
|
results: pd.DataFrame = mlflow.load_table(
|
|
"eval_results_table.json", extra_columns=["run_id", "params.prompt_template"]
|
|
)
|
|
results_grouped_by_article = results.sort_values(by="id")
|
|
print("Evaluation results:")
|
|
print(results_grouped_by_article[["run_id", "params.prompt_template", "article", "outputs"]])
|
|
|
|
# Score the best model on a new article
|
|
new_article = """
|
|
Adnan Januzaj swapped the lush turf of Old Trafford for the green baize at Sheffield when he
|
|
turned up at the snooker World Championships on Wednesday. The Manchester United winger, who has
|
|
endured a frustrating season under Louis van Gaal, had turned out for the Under 21 side at Fulham
|
|
on Tuesday night amid reports he could be farmed out on loan next season. But Januzaj may want to
|
|
consider trying his hand at another sport after displaying his silky skillls on a mini pool table.
|
|
Adnan Januzaj (left) cheered on\xa0Shaun Murphy (right) at the World Championship in Sheffield.
|
|
Januzaj shows off his potting skills on a mini pool table at the Crucible on Wednesday.
|
|
The 20-year-old Belgium international was at the Crucible to cheer on his friend Shaun Murphy in
|
|
his quarter-final against Anthony McGill. The 2005 winner moved a step closer to an elusive second
|
|
title in Sheffield with a 13-8 victory, sealed with a 67 break. Three centuries in the match, and
|
|
the way he accelerated away from 6-6, showed Murphy is a man to fear, and next for him will be
|
|
Neil Robertson or Barry Hawkins. Januzaj turned out for Under 21s in the 4-1 victory at Fulham on
|
|
Tuesday night.
|
|
"""
|
|
|
|
print(
|
|
f"Scoring the model with prompt '{prompt_template_2}' on the article '{new_article[:70] + '...'}'"
|
|
)
|
|
best_model = mlflow.pyfunc.load_model(f"runs:/{mlflow.last_active_run().info.run_id}/model")
|
|
summary = best_model.predict({"article": new_article})
|
|
print(f"Summary: {summary}")
|