@ -16,15 +16,24 @@ class PredictionGuard(LLM):
""" Wrapper around Prediction Guard large language models.
""" Wrapper around Prediction Guard large language models.
To use , you should have the ` ` predictionguard ` ` python package installed , and the
To use , you should have the ` ` predictionguard ` ` python package installed , and the
environment variable ` ` PREDICTIONGUARD_TOKEN ` ` set with your access token , or pass
environment variable ` ` PREDICTIONGUARD_TOKEN ` ` set with your access token , or pass
it as a named parameter to the constructor .
it as a named parameter to the constructor . To use Prediction Guard ' s API along
with OpenAI models , set the environment variable ` ` OPENAI_API_KEY ` ` with your
OpenAI API key as well .
Example :
Example :
. . code - block : : python
. . code - block : : python
pgllm = PredictionGuard ( name = " text-gen-proxy-name " , token = " my-access-token " )
pgllm = PredictionGuard ( model = " MPT-7B-Instruct " ,
token = " my-access-token " ,
output = {
" type " : " boolean "
} )
"""
"""
client : Any #: :meta private:
client : Any #: :meta private:
name : Optional [ str ] = " default-text-gen "
model : Optional [ str ] = " MPT-7B-Instruct "
""" Proxy name to use. """
""" Model name to use. """
output : Optional [ Dict [ str , Any ] ] = None
""" The output type or structure for controlling the LLM output. """
max_tokens : int = 256
max_tokens : int = 256
""" Denotes the number of tokens to predict per generation. """
""" Denotes the number of tokens to predict per generation. """
@ -33,6 +42,7 @@ class PredictionGuard(LLM):
""" A non-negative float that tunes the degree of randomness in generation. """
""" A non-negative float that tunes the degree of randomness in generation. """
token : Optional [ str ] = None
token : Optional [ str ] = None
""" Your Prediction Guard access token. """
stop : Optional [ List [ str ] ] = None
stop : Optional [ List [ str ] ] = None
@ -58,7 +68,7 @@ class PredictionGuard(LLM):
@property
@property
def _default_params ( self ) - > Dict [ str , Any ] :
def _default_params ( self ) - > Dict [ str , Any ] :
""" Get the default parameters for calling Cohere API."""
""" Get the default parameters for calling the Prediction Guard API."""
return {
return {
" max_tokens " : self . max_tokens ,
" max_tokens " : self . max_tokens ,
" temperature " : self . temperature ,
" temperature " : self . temperature ,
@ -67,7 +77,7 @@ class PredictionGuard(LLM):
@property
@property
def _identifying_params ( self ) - > Dict [ str , Any ] :
def _identifying_params ( self ) - > Dict [ str , Any ] :
""" Get the identifying parameters. """
""" Get the identifying parameters. """
return { * * { " na me" : self . na me} , * * self . _default_params }
return { * * { " mod el " : self . mod el } , * * self . _default_params }
@property
@property
def _llm_type ( self ) - > str :
def _llm_type ( self ) - > str :
@ -80,7 +90,7 @@ class PredictionGuard(LLM):
stop : Optional [ List [ str ] ] = None ,
stop : Optional [ List [ str ] ] = None ,
run_manager : Optional [ CallbackManagerForLLMRun ] = None ,
run_manager : Optional [ CallbackManagerForLLMRun ] = None ,
) - > str :
) - > str :
""" Call out to Prediction Guard ' s model proxy .
""" Call out to Prediction Guard ' s model API .
Args :
Args :
prompt : The prompt to pass into the model .
prompt : The prompt to pass into the model .
Returns :
Returns :
@ -89,6 +99,8 @@ class PredictionGuard(LLM):
. . code - block : : python
. . code - block : : python
response = pgllm ( " Tell me a joke. " )
response = pgllm ( " Tell me a joke. " )
"""
"""
import predictionguard as pg
params = self . _default_params
params = self . _default_params
if self . stop is not None and stop is not None :
if self . stop is not None and stop is not None :
raise ValueError ( " `stop` found in both the input and default params. " )
raise ValueError ( " `stop` found in both the input and default params. " )
@ -97,15 +109,14 @@ class PredictionGuard(LLM):
else :
else :
params [ " stop_sequences " ] = stop
params [ " stop_sequences " ] = stop
response = self . client . predict (
response = pg . Completion . create (
name = self . name ,
model = self . model ,
data = {
prompt = prompt ,
" prompt " : prompt ,
output = self . output ,
" max_tokens " : params [ " max_tokens " ] ,
temperature = params [ " temperature " ] ,
" temperature " : params [ " temperature " ] ,
max_tokens = params [ " max_tokens " ] ,
} ,
)
)
text = response [ " text" ]
text = response [ " choices" ] [ 0 ] [ " text" ]
# If stop tokens are provided, Prediction Guard's endpoint returns them.
# If stop tokens are provided, Prediction Guard's endpoint returns them.
# In order to make this consistent with other endpoints, we strip them.
# In order to make this consistent with other endpoints, we strip them.