mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-02 01:28:31 +00:00
[devops] remove post commit ci (#5566)
* [devops] remove post commit ci * [misc] run pre-commit on all files * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -21,7 +21,7 @@ print(resp) # super-heavyweight awesome-natured yawning Australian creature!
|
||||
|
||||
"""
|
||||
import json
|
||||
from typing import Any, List, Mapping, Optional
|
||||
from typing import Any, Mapping
|
||||
|
||||
import requests
|
||||
from langchain.llms.base import LLM
|
||||
@@ -31,31 +31,31 @@ from langchain.utils import get_from_dict_or_env
|
||||
class ColossalCloudLLM(LLM):
|
||||
"""
|
||||
A custom LLM class that integrates LLMs running on the ColossalCloud Platform
|
||||
|
||||
|
||||
"""
|
||||
n: int
|
||||
gen_config: dict = None
|
||||
|
||||
n: int
|
||||
gen_config: dict = None
|
||||
auth_config: dict = None
|
||||
valid_gen_para: list = ['max_new_tokens', 'top_k',
|
||||
'top_p', 'temperature', 'repetition_penalty']
|
||||
valid_gen_para: list = ["max_new_tokens", "top_k", "top_p", "temperature", "repetition_penalty"]
|
||||
|
||||
def __init__(self, gen_config=None, **kwargs):
|
||||
"""
|
||||
Args:
|
||||
gen_config: config for generation,
|
||||
max_new_tokens: 50 by default
|
||||
top_k: (1, vocab_size)
|
||||
top_k: (1, vocab_size)
|
||||
top_p: (0, 1) if not None
|
||||
temperature: (0, inf) if not None
|
||||
temperature: (0, inf) if not None
|
||||
repetition_penalty: (1, inf) if not None
|
||||
"""
|
||||
super(ColossalCloudLLM, self).__init__(**kwargs)
|
||||
if gen_config is None:
|
||||
self.gen_config = {"max_new_tokens": 50}
|
||||
else:
|
||||
if gen_config is None:
|
||||
self.gen_config = {"max_new_tokens": 50}
|
||||
else:
|
||||
assert "max_new_tokens" in gen_config, "max_new_tokens is a compulsory key in the gen config"
|
||||
self.gen_config = gen_config
|
||||
|
||||
|
||||
@property
|
||||
def _identifying_params(self) -> Mapping[str, Any]:
|
||||
"""Get the identifying parameters."""
|
||||
@@ -63,17 +63,17 @@ class ColossalCloudLLM(LLM):
|
||||
|
||||
@property
|
||||
def _llm_type(self) -> str:
|
||||
return 'ColossalCloudLLM'
|
||||
|
||||
return "ColossalCloudLLM"
|
||||
|
||||
def set_auth_config(self, **kwargs):
|
||||
url = get_from_dict_or_env(kwargs, "url", "URL")
|
||||
host = get_from_dict_or_env(kwargs, "host", "HOST")
|
||||
|
||||
|
||||
auth_config = {}
|
||||
auth_config['endpoint'] = url
|
||||
auth_config['Host'] = host
|
||||
auth_config["endpoint"] = url
|
||||
auth_config["Host"] = host
|
||||
self.auth_config = auth_config
|
||||
|
||||
|
||||
def _call(self, prompt: str, stop=None, **kwargs: Any) -> str:
|
||||
"""
|
||||
Args:
|
||||
@@ -81,15 +81,17 @@ class ColossalCloudLLM(LLM):
|
||||
stop: A list of strings to stop generation when encountered
|
||||
|
||||
Returns:
|
||||
The string generated by the model
|
||||
The string generated by the model
|
||||
"""
|
||||
# Update the generation arguments
|
||||
for key, value in kwargs.items():
|
||||
if key not in self.valid_gen_para:
|
||||
raise KeyError(f"Invalid generation parameter: '{key}'. Valid keys are: {', '.join(self.valid_gen_para)}")
|
||||
raise KeyError(
|
||||
f"Invalid generation parameter: '{key}'. Valid keys are: {', '.join(self.valid_gen_para)}"
|
||||
)
|
||||
if key in self.gen_config:
|
||||
self.gen_config[key] = value
|
||||
|
||||
|
||||
resp_text = self.text_completion(prompt, self.gen_config, self.auth_config)
|
||||
# TODO: This may cause excessive tokens count
|
||||
if stop is not None:
|
||||
@@ -97,29 +99,19 @@ class ColossalCloudLLM(LLM):
|
||||
if stopping_words in resp_text:
|
||||
resp_text = resp_text.split(stopping_words)[0]
|
||||
return resp_text
|
||||
|
||||
|
||||
def text_completion(self, prompt, gen_config, auth_config):
|
||||
# Required Parameters
|
||||
endpoint = auth_config.pop('endpoint')
|
||||
max_new_tokens = gen_config.pop('max_new_tokens')
|
||||
endpoint = auth_config.pop("endpoint")
|
||||
max_new_tokens = gen_config.pop("max_new_tokens")
|
||||
# Optional Parameters
|
||||
optional_params = ['top_k', 'top_p', 'temperature', 'repetition_penalty'] # Self.optional
|
||||
optional_params = ["top_k", "top_p", "temperature", "repetition_penalty"] # Self.optional
|
||||
gen_config = {key: gen_config[key] for key in optional_params if key in gen_config}
|
||||
# Define the data payload
|
||||
data = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"history": [
|
||||
{"instruction": prompt, "response": ""}
|
||||
],
|
||||
**gen_config
|
||||
}
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
**auth_config # 'Host',
|
||||
}
|
||||
data = {"max_new_tokens": max_new_tokens, "history": [{"instruction": prompt, "response": ""}], **gen_config}
|
||||
headers = {"Content-Type": "application/json", **auth_config} # 'Host',
|
||||
# Make the POST request
|
||||
response = requests.post(endpoint, headers=headers, data=json.dumps(data))
|
||||
response.raise_for_status() # raise error if return code is not 200(success)
|
||||
response.raise_for_status() # raise error if return code is not 200(success)
|
||||
# Check the response
|
||||
return response.text
|
||||
|
Reference in New Issue
Block a user