logo
Llama-2
repo-copy-icon

copied

Browse Source

Replace with chat completion

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 1 year ago
parent
commit
fd410c19e2
  1. 2
      README.md
  2. 26
      llama2.py

2
README.md

@ -25,7 +25,7 @@ Use the default model to continue the conversation from given messages.
```python ```python
from towhee import ops from towhee import ops
chat = ops.LLM.Llama_2('llama-2-13b-chat', max_tokens=512)
chat = ops.LLM.Llama_2('llama-2-13b-chat', n_ctx=4096, max_tokens=200)
message = [ message = [
{'system': 'You are a very helpful assistant.'}, {'system': 'You are a very helpful assistant.'},

26
llama2.py

@ -39,36 +39,32 @@ class LlamaCpp(PyOperator):
self.model_path = model_name_or_file self.model_path = model_name_or_file
assert os.path.isfile(self.model_path), f'Invalid model path: {self.model_path}' assert os.path.isfile(self.model_path), f'Invalid model path: {self.model_path}'
self.model = Llama(model_path=self.model_path)
init_kwargs = {}
for k in vars(Llama).keys() & self.kwargs.keys():
init_kwargs[k] = self.kwargs.pop(k)
self.model = Llama(model_path=self.model_path, **init_kwargs)
def __call__(self, messages: List[dict]): def __call__(self, messages: List[dict]):
prompt = self.parse_inputs(messages)
resp = self.model(prompt, **self.kwargs)
messages = self.parse_inputs(messages)
resp = self.model.create_chat_completion(messages, **self.kwargs)
answer = self.parse_outputs(resp) answer = self.parse_outputs(resp)
return answer return answer
def parse_inputs(self, messages: List[dict]): def parse_inputs(self, messages: List[dict]):
assert isinstance(messages, list), \ assert isinstance(messages, list), \
'Inputs must be a list of dictionaries with keys from ["system", "question", "answer"].' 'Inputs must be a list of dictionaries with keys from ["system", "question", "answer"].'
prompt = ''
question = messages.pop(-1)
assert len(question) == 1 and 'question' in question.keys()
question = question['question']
new_messages = []
for m in messages: for m in messages:
for k, v in m.items(): for k, v in m.items():
if k == 'system': if k == 'system':
prompt += f'''[INST] <<SYS>> {v} <</SYS>>\n'''
new_messages.append({'role': 'system', 'content': v})
elif k == 'question': elif k == 'question':
prompt += f''' {v} [/INST]\n'''
new_messages.append({'role': 'user', 'content': v})
elif k == 'answer': elif k == 'answer':
prompt += f''' {v} '''
new_messages.append({'role': 'assistant', 'content': v})
else: else:
raise KeyError(f'Invalid key of message: {k}') raise KeyError(f'Invalid key of message: {k}')
if len(prompt) > 0:
prompt = '<s> ' + prompt + ' </s>' + f'<s> [INST] {question} [/INST]'
else:
prompt = question
return prompt
return new_messages
def parse_outputs(self, response): def parse_outputs(self, response):
return response['choices'][0]['text'] return response['choices'][0]['text']

Loading…
Cancel
Save