stark_qa.tools.api_lib

stark_qa.tools.api_lib.claude

stark_qa.tools.api_lib.claude.complete_text_claude(message, model='claude-2.1', json_object=False, max_tokens=2048, temperature=1.0, max_retry=1, sleep_time=0, tools=[], **kwargs)[source]

Call the Claude API to complete a prompt.

Parameters:

  • message (Union[str, list]) – The input message or a list of message dicts.

  • model (str) – The model to use for completion. Default is “claude-2.1”.

  • json_object (bool) – Whether to output in JSON format. Default is False.

  • max_tokens (int) – Maximum number of tokens to generate. Default is 2048.

  • temperature (float) – Sampling temperature. Default is 1.0.

  • max_retry (int) – Maximum number of retries in case of an error. Default is 1.

  • sleep_time (int) – Sleep time between retries in seconds. Default is 0.

  • tools (list) – List of tools to use for the completion. Default is an empty list.

  • **kwargs (Any) – Additional keyword arguments to pass to the API.

Returns:

The completed text generated by the Claude model.

Return type:

str

Raises:

Exception – If the completion fails after the maximum number of retries.

stark_qa.tools.api_lib.gpt

stark_qa.tools.api_lib.gpt.get_gpt_output(message, model='gpt-4-1106-preview', max_tokens=2048, temperature=1.0, max_retry=1, sleep_time=60, json_object=False)[source]

Call the OpenAI API to get the GPT model output for a given prompt.

Parameters:

  • message (Union[str, List[Dict[str, str]]]) – The input message or a list of message dicts.

  • model (str) – The model to use for completion. Default is “gpt-4-1106-preview”.

  • max_tokens (int) – Maximum number of tokens to generate. Default is 2048.

  • temperature (float) – Sampling temperature. Default is 1.0.

  • max_retry (int) – Maximum number of retries in case of an error. Default is 1.

  • sleep_time (int) – Sleep time between retries in seconds. Default is 60.

  • json_object (bool) – Whether to output in JSON format. Default is False.

Returns:

The completed text generated by the GPT model.

Return type:

str

Raises:

Exception – If the completion fails after the maximum number of retries.

stark_qa.tools.api_lib.huggingface

stark_qa.tools.api_lib.huggingface.complete_text_hf(message, model='huggingface/codellama/CodeLlama-7b-hf', max_tokens=2000, temperature=0.5, json_object=False, max_retry=1, sleep_time=0, stop_sequences=[], **kwargs)[source]

Generate text completion using a specified Hugging Face model.

Parameters:

  • message (str) – The input text message for completion.

  • model (str) – The Hugging Face model to use. Default is “huggingface/codellama/CodeLlama-7b-hf”.

  • max_tokens (int) – The maximum number of tokens to generate. Default is 2000.

  • temperature (float) – Sampling temperature for generation. Default is 0.5.

  • json_object (bool) – Whether to format the message for JSON output. Default is False.

  • max_retry (int) – Maximum number of retries in case of an error. Default is 1.

  • sleep_time (int) – Sleep time between retries in seconds. Default is 0.

  • stop_sequences (list) – List of stop sequences to halt the generation.

  • **kwargs – Additional keyword arguments for the generate function.

Returns:

The generated text completion.

Return type:

str

stark_qa.tools.api_lib.openai_emb

stark_qa.tools.api_lib.openai_emb.get_openai_embedding(text, model='text-embedding-ada-002', max_retry=1, sleep_time=0)[source]

Get the OpenAI embedding for a given text.

Parameters:

  • text (str) – The input text to be embedded.

  • model (str) – The model to use for embedding. Default is “text-embedding-ada-002”.

  • max_retry (int) – Maximum number of retries in case of an error. Default is 1.

  • sleep_time (int) – Sleep time between retries in seconds. Default is 0.

Returns:

The embedding of the input text.

Return type:

torch.FloatTensor

stark_qa.tools.api_lib.openai_emb.get_openai_embeddings(texts, n_max_nodes=5, model='text-embedding-ada-002')[source]

Get embeddings for a list of texts using OpenAI’s embedding model.

Parameters:

  • texts (list) – List of input texts to be embedded.

  • n_max_nodes (int) – Maximum number of parallel processes. Default is 5.

  • model (str) – The model to use for embedding. Default is “text-embedding-ada-002”.

Returns:

A tensor containing embeddings for all input texts.

Return type:

torch.FloatTensor