github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-05 20:58:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Harrison/handle stop tokens ai21 (#1077)

Browse Source 
Co-authored-by: Andrew Huang <jhuang16888@gmail.com>
This commit is contained in:
Harrison Chase 2023-02-15 22:44:55 -08:00 committed by GitHub
parent d8ed286200
commit 52753066ef
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
langchain/llms/ai21.py
View File

@ -1,5 +1,5 @@
"""Wrapper around AI21 APIs."""
from typing import Any, Dict, List, Mapping, Optional
from typing import Any, Dict, List, Optional
import requests
from pydantic import BaseModel, Extra, root_validator
@ -64,6 +64,8 @@ class AI21(LLM, BaseModel):
ai21_api_key: Optional[str] = None
stop: Optional[List[str]] = None
base_url: Optional[str] = None
"""Base url to use, if None decides based on model name."""
@ -80,7 +82,7 @@ class AI21(LLM, BaseModel):
return values
@property
def _default_params(self) -> Mapping[str, Any]:
def _default_params(self) -> Dict[str, Any]:
"""Get the default parameters for calling AI21 API."""
return {
"temperature": self.temperature,
@ -95,7 +97,7 @@ class AI21(LLM, BaseModel):
}
@property
def _identifying_params(self) -> Mapping[str, Any]:
def _identifying_params(self) -> Dict[str, Any]:
"""Get the identifying parameters."""
return {**{"model": self.model}, **self._default_params}
@ -119,7 +121,11 @@ class AI21(LLM, BaseModel):
response = ai21("Tell me a joke.")
"""
if stop is None:
if self.stop is not None and stop is not None:
raise ValueError("`stop` found in both the input and default params.")
elif self.stop is not None:
stop = self.stop
elif stop is None:
stop = []
if self.base_url is not None:
base_url = self.base_url