Skip to content

Add vllm support #127

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
barry-ravichandran wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add vllm support #127

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 33 additions & 17 deletions align_system/algorithms/outlines_adm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import outlines
from outlines.samplers import MultinomialSampler
from outlines.models.vllm import VLLM
import jinja2
from rich.highlighter import JSONHighlighter
from swagger_client.models import (
Expand Down Expand Up @@ -73,11 +74,10 @@ def __init__(self,

model_kwargs['torch_dtype'] = torch_dtype

self.model = outlines.models.transformers(
self.model = outlines.models.vllm(
model_name,
device=device,
model_kwargs=model_kwargs,
tokenizer_kwargs=kwargs.get('tokenizer_kwargs', {}))
device=device
)
# NOTE: In cases where we want multiple samples, we're passing
# in a list of prompts (this allows us to shuffle answers in
# each prompt), rather than setting the number of samples in
Expand All @@ -86,7 +86,7 @@ def __init__(self,
self.sampler = sampler

def dialog_to_prompt(self, dialog):
tokenizer = self.model.tokenizer.tokenizer
tokenizer = self.model.tokenizer

try:
encoded_dialog = tokenizer.apply_chat_template(dialog)
Expand Down Expand Up @@ -218,32 +218,48 @@ def top_level_choose_action(self,
if "incontext" in kwargs and "number" in incontext_settings and incontext_settings["number"] > 0:
scenario_to_match = scenario_state_description_1(scenario_state)
prompt_to_match, _ = self._state_to_top_level_prompt(scenario_state, available_actions)

# Create positive ICL example generators
positive_target = {'kdma': kdma, 'name': name, 'value': value}
positive_icl_example_generator = incontext_utils.BaselineIncontextExampleGenerator(incontext_settings,
[positive_target])
positive_icl_example_generator = (
incontext_utils.BaselineIncontextExampleGenerator(
incontext_settings,
[positive_target]
)
)
# Get subset of relevant of examples
positive_selected_icl_examples = positive_icl_example_generator.select_icl_examples(kdma,
scenario_to_match,
prompt_to_match)
positive_selected_icl_examples = (
positive_icl_example_generator.select_icl_examples(
kdma,
scenario_to_match,
prompt_to_match
)
)
# Create positive ICL prompts
for icl_sample in positive_selected_icl_examples:
positive_icl_examples.extend([
{"role": "user", "content": icl_sample['prompt']},
{"role": "assistant", "content": f'{icl_sample["response"]}'}
])

# Create negative ICL prompts
if num_negative_samples > 0:
# Create negative ICL example generators
negative_target = {'kdma': kdma, 'name': name, 'value': negative_value}
negative_icl_example_generator = incontext_utils.BaselineIncontextExampleGenerator(incontext_settings,
[negative_target])
negative_icl_example_generator = (
incontext_utils.BaselineIncontextExampleGenerator(
incontext_settings,
[negative_target]
)
)
# Get subset of relevant of examples
negative_selected_icl_examples = negative_icl_example_generator.select_icl_examples(kdma,
scenario_to_match,
prompt_to_match)
negative_selected_icl_examples = (
negative_icl_example_generator.select_icl_examples(
kdma,
scenario_to_match,
prompt_to_match
)
)
for icl_sample in negative_selected_icl_examples:
negative_icl_examples.extend([
{"role": "user", "content": icl_sample['prompt']},
Expand Down
Loading