Spaces:
Runtime error
Runtime error
File size: 3,051 Bytes
ed4d993 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
from __future__ import annotations
import json
import logging
from typing import Any, Dict, List, Optional
import requests
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.pydantic_v1 import Field, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
from langchain_community.llms.utils import enforce_stop_tokens
logger = logging.getLogger(__name__)
class BaichuanLLM(LLM):
# TODO: Adding streaming support.
"""Baichuan large language models."""
model: str = "Baichuan2-Turbo-192k"
"""
Other models are available at https://platform.baichuan-ai.com/docs/api.
"""
temperature: float = 0.3
top_p: float = 0.95
timeout: int = 60
model_kwargs: Dict[str, Any] = Field(default_factory=dict)
baichuan_api_host: Optional[str] = None
baichuan_api_key: Optional[SecretStr] = None
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
values["baichuan_api_key"] = convert_to_secret_str(
get_from_dict_or_env(values, "baichuan_api_key", "BAICHUAN_API_KEY")
)
values["baichuan_api_host"] = get_from_dict_or_env(
values,
"baichuan_api_host",
"BAICHUAN_API_HOST",
default="https://api.baichuan-ai.com/v1/chat/completions",
)
return values
@property
def _default_params(self) -> Dict[str, Any]:
return {
"model": self.model,
"temperature": self.temperature,
"top_p": self.top_p,
**self.model_kwargs,
}
def _post(self, request: Any) -> Any:
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {self.baichuan_api_key.get_secret_value()}", # type: ignore[union-attr]
}
try:
response = requests.post(
self.baichuan_api_host, # type: ignore[arg-type]
headers=headers,
json=request,
timeout=self.timeout,
)
if response.status_code == 200:
parsed_json = json.loads(response.text)
return parsed_json["choices"][0]["message"]["content"]
else:
response.raise_for_status()
except Exception as e:
raise ValueError(f"An error has occurred: {e}")
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
request = self._default_params
request["messages"] = [{"role": "user", "content": prompt}]
request.update(kwargs)
text = self._post(request)
if stop is not None:
text = enforce_stop_tokens(text, stop)
return text
@property
def _llm_type(self) -> str:
"""Return type of chat_model."""
return "baichuan-llm"
|