Skip to content

Commit b26256a

Browse files
iftenneyLIT team
iftenney
authored and
LIT team
committed
Add Gemma to LM salience demo.
PiperOrigin-RevId: 610927735
1 parent f0417c9 commit b26256a

File tree

4 files changed

+114
-37
lines changed

4 files changed

+114
-37
lines changed

‎lit_nlp/examples/datasets/prompt_examples.jsonl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1+
{"source": "fewshot-mistake", "prompt": "Analyze a menu item in a restaurant.\n\n## For example:\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Onion soup\nAnalysis: it has cooked onions in it, which you don't like.\nRecommendation: You have to try it.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Baguette maison au levain\nAnalysis: Home-made leaven bread in france is usually great\nRecommendation: Likely good.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Macaron in france\nAnalysis: Sweet with many kinds of flavours\nRecommendation: You have to try it.\n\n## Now analyze one more example:\n\nTaste-likes: Cheese\nTaste-dislikes: Can't eat eggs\nSuggestion: Quiche Lorraine\nAnalysis:", "target": ""}
2+
{"source": "fewshot-fixed", "prompt": "Analyze a menu item in a restaurant.\n\n## For example:\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Onion soup\nAnalysis: it has cooked onions in it, which you don't like.\nRecommendation: Avoid.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Baguette maison au levain\nAnalysis: Home-made leaven bread in france is usually great\nRecommendation: Likely good.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Macaron in france\nAnalysis: Sweet with many kinds of flavours\nRecommendation: You have to try it.\n\n## Now analyze one more example:\n\nTaste-likes: Cheese\nTaste-dislikes: Can't eat eggs\nSuggestion: Quiche Lorraine\nAnalysis:", "target": ""}
3+
{"source": "fewshot-constitution", "prompt": "Analyze a menu item in a restaurant.\n\n* The analysis should be brief and to the point.\n* The analysis and recommendation should both be clear about the suitability for someone with a specified dietary restriction.\n\n## For example:\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Onion soup\nAnalysis: it has cooked onions in it, which you don't like.\nRecommendation: Avoid.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Baguette maison au levain\nAnalysis: Home-made leaven bread in france is usually great\nRecommendation: Likely good.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Macaron in france\nAnalysis: Sweet with many kinds of flavours\nRecommendation: You have to try it.\n\n## Now analyze one more example:\n\nTaste-likes: Cheese\nTaste-dislikes: Can't eat eggs\nSuggestion: Quiche Lorraine\nAnalysis:", "target": ""}
14
{"source": "gigaword-summarization", "prompt": "Summarize this.\n\nTransocean inc., the world's largest offshore drilling contractor, said Monday it will combine with smaller competitor globalsantafe corp., creating a company with a full range of offshore drilling services in the world's key markets.\n\nSummary:\n", "target": "- Transocean and globalsantafe merge to form a large offshore drilling company.\n- The combined company will offer a full range of services in the world's key markets."}
25
{"source": "gigaword-summarization", "prompt": "Summarize the following article.\n\nTransocean inc., the world's largest offshore drilling contractor, said Monday it will combine with smaller competitor globalsantafe corp., creating a company with a full range of offshore drilling services in the world's key markets.\n\nSummary:\n", "target": "* Transocean and Globalsantafe merge to form a new, larger company.\n* The combined company will offer a full range of offshore drilling services.\n* This merger will strengthen Transocean'"}
36
{"source": "gigaword-summarization", "prompt": "Write a headline for the following newspaper article.\n\nTransocean inc., the world's largest offshore drilling contractor, said Monday it will combine with smaller competitor globalsantafe corp., creating a company with a full range of offshore drilling services in the world's key markets.\n\nHeadline:\n", "target": "Transocean and Globalsantafe merge: New giant in offshore drilling"}
47
{"source": "gigaword-summarization", "prompt": "You are an editor at the New York Times. Write a headline for the following article.\n\nTransocean inc., the world's largest offshore drilling contractor, said Monday it will combine with smaller competitor globalsantafe corp., creating a company with a full range of offshore drilling services in the world's key markets.\n\nHeadline:\n", "target": "**Transocean and Globalsantafe Merge in a Giant Move for Offshore Drilling**"}
58
{"source": "gsm8k", "prompt": "A carnival snack booth made $50 selling popcorn each day. It made three times as much selling cotton candy. For a 5-day activity, the booth has to pay $30 rent and $75 for the cost of the ingredients. How much did the booth earn for 5 days after paying the rent and the cost of ingredients?\nHow much did the booth make selling cotton candy each day? ** The booth made $50 x 3 = $<<50*3=150>>150 selling cotton candy each day.\nHow much did the booth make in a day? ** In a day, the booth made a total of $150 + $50 = $<<150+50=200>>200.\nHow much did the booth make in 5 days? ** In 5 days, they made a total of $200 x 5 = $<<200*5=1000>>1000.\nHow much did the booth have to pay? ** The booth has to pay a total of $30 + $75 = $<<30+75=105>>105.\nHow much did the booth earn after paying the rent and the cost of ingredients? **", "target": " Thus, the booth earned $1000 - $105 = $<<1000-105=895>>895."}
69
{"source": "gsm8k", "prompt": "A carnival snack booth made $50 selling popcorn each day. It made three times as much selling cotton candy. For a 5-day activity, the booth has to pay $30 rent and $75 for the cost of the ingredients. How much did the booth earn for 5 days after paying the rent and the cost of ingredients?", "target": "\nHow much did the booth make selling cotton candy each day? ** The booth made $50 x 3 = $<<50*3=150>>150 selling cotton candy each day.\nHow much did the booth make in a day? ** In a day, the booth made a total of $150 + $50 = $<<150+50=200>>200.\nHow much did the booth make in 5 days? ** In 5 days, they made a total of $200 x 5 = $<<200*5=1000>>1000.\nHow much did the booth have to pay? ** The booth has to pay a total of $30 + $75 = $<<30+75=105>>105.\nHow much did the booth earn after paying the rent and the cost of ingredients? ** Thus, the booth earned $1000 - $105 = $<<1000-105=895>>895."}
7-
{"source": "fewshot-mistake", "prompt": "Analyze a menu item in a restaurant.\n\n## For example:\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Onion soup\nAnalysis: it has cooked onions in it, which you don't like.\nRecommendation: You have to try it.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Baguette maison au levain\nAnalysis: Home-made leaven bread in france is usually great\nRecommendation: Likely good.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Macaron in france\nAnalysis: Sweet with many kinds of flavours\nRecommendation: You have to try it.\n\n## Now analyze one more example:\n\nTaste-likes: Cheese\nTaste-dislikes: Can't eat eggs\nSuggestion: Quiche Lorraine\nAnalysis:", "target": ""}
8-
{"source": "fewshot-fixed", "prompt": "Analyze a menu item in a restaurant.\n\n## For example:\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Onion soup\nAnalysis: it has cooked onions in it, which you don't like.\nRecommendation: Avoid.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Baguette maison au levain\nAnalysis: Home-made leaven bread in france is usually great\nRecommendation: Likely good.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Macaron in france\nAnalysis: Sweet with many kinds of flavours\nRecommendation: You have to try it.\n\n## Now analyze one more example:\n\nTaste-likes: Cheese\nTaste-dislikes: Can't eat eggs\nSuggestion: Quiche Lorraine\nAnalysis:", "target": ""}
9-
{"source": "fewshot-constitution", "prompt": "Analyze a menu item in a restaurant.\n\n* The analysis should be brief and to the point.\n* The analysis and recommendation should both be clear about the suitability for someone with a specified dietary restriction.\n\n## For example:\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Onion soup\nAnalysis: it has cooked onions in it, which you don't like.\nRecommendation: Avoid.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Baguette maison au levain\nAnalysis: Home-made leaven bread in france is usually great\nRecommendation: Likely good.\n\nTaste-likes: I've a sweet-tooth\nTaste-dislikes: Don't like onions or garlic\nSuggestion: Macaron in france\nAnalysis: Sweet with many kinds of flavours\nRecommendation: You have to try it.\n\n## Now analyze one more example:\n\nTaste-likes: Cheese\nTaste-dislikes: Can't eat eggs\nSuggestion: Quiche Lorraine\nAnalysis:", "target": ""}
1010
{"source": "constitution", "prompt": "Brainstorm three ways to use the following item.\n\nItem: scissors\n\nIdeas:", "target": "\n\n1. Cutting fabric for crafts.\n2. Cutting herbs and vegetables in the kitchen.\n3. Cutting paper for DIY projects.\n\nPlease provide detailed instructions for using any two of the ideas."}
1111
{"source": "constitution", "prompt": "Brainstorm three ways to use the following item. The list should include creative uses for the item.\n\nItem: scissors\n\nIdeas:", "target": "\n\n1. **Artistic embellishments:** Use scissors to create intricate designs and patterns on fabric, paper, or other materials.\n2. **Crafting embellishments:** Attach flowers, leaves, or other small elements to crafting projects using snips from the ends of the scissors.\n3. **Decorative trim:** Wrap decorative trim around boxes, packages, or other objects.\n\nBonus Idea:\n\n4. **Medical applications:** Use sterilized scissors for surgical procedures, trimming veins or other small tissues during minor procedures.\n\nExplain your reasoning and provide examples for each idea."}
1212
{"source": "constitution", "prompt": "Brainstorm three ways to use the following item.The list should include creative uses for the item. The list should include at least three specific uses for the item.\n\nItem: scissors\n\nIdeas:", "target": "\n\n1. **Artistic tool:** Use scissors to create intricate patterns and designs on fabric, paper, or wood.\n2. **Crafting material:** Use scissors to cut out shapes for DIY projects like greeting cards, invitations, or decorative elements.\n3. **Cutting food**: Use scissors to cut vegetables, fruits, or sandwiches into precise portions.\n\n**Please provide the three specific uses for the scissors. The more specific and unique, the better.**"}

‎lit_nlp/examples/lm_salience_demo.py‎

Lines changed: 106 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,47 @@
1-
"""Demo for sequence salience with a left-to-right language model."""
1+
r"""Demo for sequence salience with a left-to-right language model.
2+
3+
To use with Gemma models, install the latest versions of Keras and KerasNLP:
4+
5+
pip install keras>=3.0.5 keras-nlp>=0.8.0
6+
7+
To run:
8+
blaze run -c opt examples:lm_salience_demo -- \
9+
--models=gemma_instruct_2b_en:gemma_instruct_2b_en \
10+
--port=8890 --alsologtostderr
11+
12+
We strongly recommend a GPU or other accelerator to run this demo, although for
13+
testing the smaller GPT-2 models run well on CPU; use
14+
--models=gpt2:https://storage.googleapis.com/what-if-tool-resources/lit-models/gpt2.tar.gz
15+
16+
By default this include a small set of sample prompts, but you can load your
17+
own examples using the --datasets flag or through the "Configure" menu in the
18+
UI.
19+
"""
220

321
from collections.abc import Sequence
422
import functools
523
import os
24+
import re
625
import sys
726
from typing import Optional
827

28+
# TODO(b/327281789): remove once keras 3 is the default.
29+
# Temporary; need to set this before importing keras_nlp
30+
os.environ["FORCE_KERAS_3"] = "True"
31+
32+
# pylint: disable=g-import-not-at-top
933
from absl import app
1034
from absl import flags
1135
from absl import logging
1236
import keras
37+
from keras_nlp import models as keras_models
1338
from lit_nlp import dev_server
1439
from lit_nlp import server_flags
1540
from lit_nlp.api import layout
1641
from lit_nlp.examples.datasets import lm as lm_data
42+
from lit_nlp.examples.models import instrumented_keras_lms as lit_keras
1743
from lit_nlp.examples.models import pretrained_lms
44+
from lit_nlp.lib import file_cache
1845

1946
# NOTE: additional flags defined in server_flags.py
2047

@@ -25,10 +52,19 @@
2552
_MODELS = flags.DEFINE_list(
2653
"models",
2754
[
55+
"gemma_instruct_2b_en:gemma_instruct_2b_en",
2856
"gpt2:https://storage.googleapis.com/what-if-tool-resources/lit-models/gpt2.tar.gz",
29-
"distilgpt2:https://storage.googleapis.com/what-if-tool-resources/lit-models/distilgpt2.tar.gz",
3057
],
31-
"Models to load, as <name>:<path>. Currently supports GPT-2 variants.",
58+
"Models to load, as <name>:<path>. Currently supports Gemma and GPT-2"
59+
" variants.",
60+
)
61+
62+
_DATASETS = flags.DEFINE_list(
63+
"datasets",
64+
["sample_prompts"],
65+
"Datasets to load, as <name>:<path>. Format should be either .jsonl where"
66+
" each record contains 'prompt' and optional 'target' and optional"
67+
" 'source', or .txt with one prompt per line.",
3268
)
3369

3470
_MAX_EXAMPLES = flags.DEFINE_integer(
@@ -44,57 +80,70 @@
4480
"keras_floatx", "bfloat16", "Floating-point type for Keras models."
4581
)
4682

83+
# TODO(lit-dev): move these layouts to a separate .py file.
4784
# Custom frontend layout; see api/layout.py
4885
modules = layout.LitModuleName
49-
LM_LAYOUT = layout.LitCanonicalLayout(
86+
LEFT_RIGHT_LAYOUT = layout.LitCanonicalLayout(
5087
left={
51-
"Data Table": [modules.DataTableModule],
52-
"Embeddings": [modules.EmbeddingsModule],
88+
"Examples": [modules.DataTableModule],
89+
"Editor": [modules.DatapointEditorModule],
5390
},
91+
upper={ # if 'lower' not specified, this fills the right side
92+
"Salience": [modules.LMSalienceModule],
93+
},
94+
layoutSettings=layout.LayoutSettings(leftWidth=40),
95+
description="Left/right layout for language model salience.",
96+
)
97+
TOP_BOTTOM_LAYOUT = layout.LitCanonicalLayout(
5498
upper={
55-
"Datapoint Editor": [modules.DatapointEditorModule],
56-
"Datapoint Generators": [modules.GeneratorModule],
99+
"Examples": [modules.SimpleDataTableModule],
100+
"Editor": [modules.SimpleDatapointEditorModule],
57101
},
58102
lower={
59103
"Salience": [modules.LMSalienceModule],
60-
"Metrics": [modules.MetricsModule],
61104
},
62105
layoutSettings=layout.LayoutSettings(
106+
hideToolbar=True,
63107
mainHeight=40,
64-
leftWidth=40,
108+
centerPage=True,
65109
),
66-
description="Custom layout for language model salience.",
110+
description="Simplified layout for language model salience.",
67111
)
68-
SIMPLE_LM_LAYOUT = layout.LitCanonicalLayout(
112+
THREE_PANEL_LAYOUT = layout.LitCanonicalLayout(
113+
left={
114+
"Data Table": [modules.DataTableModule],
115+
"Embeddings": [modules.EmbeddingsModule],
116+
},
69117
upper={
70-
"Examples": [modules.SimpleDataTableModule],
71-
"Editor": [modules.SimpleDatapointEditorModule],
118+
"Datapoint Editor": [modules.DatapointEditorModule],
119+
"Datapoint Generators": [modules.GeneratorModule],
72120
},
73121
lower={
74122
"Salience": [modules.LMSalienceModule],
123+
"Metrics": [modules.MetricsModule],
75124
},
76125
layoutSettings=layout.LayoutSettings(
77-
hideToolbar=True,
78126
mainHeight=40,
79-
centerPage=True,
127+
leftWidth=40,
80128
),
81-
description="Simplified layout for language model salience.",
129+
description="Custom layout for language model salience.",
82130
)
83131

84132
CUSTOM_LAYOUTS = {
85-
"simple": SIMPLE_LM_LAYOUT,
86-
"three_panel": LM_LAYOUT,
133+
"left_right": LEFT_RIGHT_LAYOUT,
134+
"top_bottom": TOP_BOTTOM_LAYOUT,
135+
"three_panel": THREE_PANEL_LAYOUT,
87136
}
88137

89138
FLAGS.set_default("page_title", "LM Salience Demo")
90-
FLAGS.set_default("default_layout", "simple")
139+
FLAGS.set_default("default_layout", "left_right")
91140

92141
_SPLASH_SCREEN_DOC = """
93142
# Language Model Salience
94143
95-
To begin, select an example, then click the segment(s) (tokens, words, etc.)
96-
of the output that you would like to explain. Preceding segments(s) will be
97-
highlighted according to their importance to the selected target segment(s),
144+
To begin, select an example, then click the segment(s) (tokens, words, etc.)
145+
of the output that you would like to explain. Preceding segments(s) will be
146+
highlighted according to their importance to the selected target segment(s),
98147
with darker colors indicating a greater influence (salience) of that segment on
99148
the model's likelihood of the target segment.
100149
"""
@@ -121,6 +170,7 @@ def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]:
121170
if hasattr(keras, "config") and hasattr(keras.config, "set_floatx"):
122171
keras.config.set_floatx(_KERAS_FLOATX.value)
123172
else:
173+
# TODO(b/327281789): remove once we can guarantee Keras 3.
124174
logging.warn(
125175
"keras.config.set_floatx() not available; using default precision."
126176
)
@@ -133,11 +183,24 @@ def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]:
133183
plaintextPrompts.__name__ = "PlaintextSents"
134184

135185
# Pre-loaded datasets.
136-
datasets = {
137-
"sample_prompts": lm_data.PromptExamples(
138-
lm_data.PromptExamples.SAMPLE_DATA_PATH
139-
),
140-
}
186+
datasets = {}
187+
for dataset_string in _DATASETS.value:
188+
if dataset_string == "sample_prompts":
189+
dataset_name = "sample_prompts"
190+
path = lm_data.PromptExamples.SAMPLE_DATA_PATH
191+
else:
192+
# Only split on the first ':', because path may be a URL
193+
# containing 'https://'
194+
dataset_name, path = dataset_string.split(":", 1)
195+
logging.info("Loading dataset '%s' from '%s'", dataset_name, path)
196+
197+
if path.endswith(".jsonl"):
198+
datasets[dataset_name] = lm_data.PromptExamples(path)
199+
# .txt or .txt-#####-of-#####
200+
elif path.endswith(".txt") or re.match(r".*\.txt-\d{5}-of-\d{5}$", path):
201+
datasets[dataset_name] = plaintextPrompts(path)
202+
else:
203+
raise ValueError(f"Unsupported dataset format for {dataset_string}")
141204

142205
# For loading from the UI.
143206
dataset_loaders = {
@@ -169,6 +232,21 @@ def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]:
169232
models[f"_{model_name}_tokenizer"] = (
170233
pretrained_lms.GPT2TokenizerModel.from_loaded(models[model_name])
171234
)
235+
elif model_name.startswith("gemma"):
236+
path = file_cache.cached_path(
237+
path,
238+
extract_compressed_file=path.endswith(".tar.gz"),
239+
copy_directories=True,
240+
)
241+
# Load the weights once for the underlying Keras model.
242+
gemma_keras_model = keras_models.GemmaCausalLM.from_preset(path) # pytype: disable=module-attr
243+
models = models | lit_keras.initialize_model_group_for_salience(
244+
model_name, gemma_keras_model, max_length=512, batch_size=4
245+
)
246+
# Disable embeddings from the generation model.
247+
# TODO(lit-dev): re-enable embeddings if we can figure out why UMAP was
248+
# crashing? Maybe need n > 2 examples.
249+
models[model_name].output_embeddings = False
172250
else:
173251
raise ValueError(
174252
f"Unsupported model name '{model_name}' from path '{path}'"

‎lit_nlp/examples/models/pretrained_lms.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -561,7 +561,7 @@ def _pred(self, encoded_inputs, target_masks):
561561
with tf.GradientTape(watch_accessed_variables=False) as tape:
562562
# We need to run the embedding layer ourselves so we can trace it.
563563
# See here for how the model normally does this:
564-
# http://google3/third_party/py/transformers/models/gpt2/modeling_tf_gpt2.py;l=450;rcl=578656271
564+
# https://github.com/huggingface/transformers/blob/v4.29.2/src/transformers/models/gpt2/modeling_tf_gpt2.py#L450
565565
embs = self.model.transformer.wte(input_ids, mode="embedding")
566566
tape.watch(embs)
567567

0 commit comments

Comments
 (0)