Commit 68c30aa7 authored by haileyschoelkopf's avatar haileyschoelkopf
Browse files

push most recent code

parent b8bda478
...@@ -96,6 +96,10 @@ class HFLM(LM): ...@@ -96,6 +96,10 @@ class HFLM(LM):
# PEFT and quantization options # PEFT and quantization options
peft: Optional[str] = None, peft: Optional[str] = None,
autogptq: Optional[Union[bool, str]] = False, autogptq: Optional[Union[bool, str]] = False,
# Chat templating settings
use_chat_template: Optional[bool] = False,
# TODO: validate a template exists in tokenizer config, if this flag is true
system_prompt: Optional[str] = None,
**kwargs, **kwargs,
) -> None: ) -> None:
super().__init__() super().__init__()
...@@ -241,6 +245,9 @@ class HFLM(LM): ...@@ -241,6 +245,9 @@ class HFLM(LM):
else: else:
self.tokenizer.add_special_tokens({"pad_token": "<|pad|>"}) self.tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
self.system_prompt = system_prompt
self.use_chat_template = use_chat_template
self._max_length = max_length self._max_length = max_length
self.batch_schedule = 1 self.batch_schedule = 1
...@@ -691,9 +698,11 @@ class HFLM(LM): ...@@ -691,9 +698,11 @@ class HFLM(LM):
context, continuation = req.args[0].strip(), req.args[1] context, continuation = req.args[0].strip(), req.args[1]
chat = [] chat = []
if self.system_prompt is not None: if self.system_prompt is not None:
chat += {"role": "system", "content": "You are a helpful assistant."} chat += [{"role": "system", "content": "You are a helpful assistant."}]
chat += ({"role": "user", "content": context},) chat += [
{"role": "user", "content": context},
]
# TODO: expose settings for chat formatting: # TODO: expose settings for chat formatting:
# - whether some "trigger" / start of assistant response might be placed in assistant's generation for it # - whether some "trigger" / start of assistant response might be placed in assistant's generation for it
# - if few-shot, should the fewshots be placed in separate convo turns? provided in user's single turn?... # - if few-shot, should the fewshots be placed in separate convo turns? provided in user's single turn?...
...@@ -786,6 +795,7 @@ class HFLM(LM): ...@@ -786,6 +795,7 @@ class HFLM(LM):
return context_enc, continuation_enc return context_enc, continuation_enc
def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]: def loglikelihood(self, requests: List[Instance]) -> List[Tuple[float, bool]]:
if self.use_chat_template:
print(f"First element before prompt formatting...\n{requests[0].args}") print(f"First element before prompt formatting...\n{requests[0].args}")
requests = self.wrap_chat_template(requests) requests = self.wrap_chat_template(requests)
print(f"First element after prompt formatting...\n{requests[0].args}") print(f"First element after prompt formatting...\n{requests[0].args}")
...@@ -1064,6 +1074,7 @@ class HFLM(LM): ...@@ -1064,6 +1074,7 @@ class HFLM(LM):
return re_ord.get_original(res) return re_ord.get_original(res)
def generate_until(self, requests: List[Instance]) -> List[str]: def generate_until(self, requests: List[Instance]) -> List[str]:
if self.use_chat_template:
print(f"First element before prompt formatting...\n{requests[0].args}") print(f"First element before prompt formatting...\n{requests[0].args}")
requests = self.tok_chat_template(requests) requests = self.tok_chat_template(requests)
print(f"First element after prompt formatting...\n{requests[0].args}") print(f"First element after prompt formatting...\n{requests[0].args}")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment