Module glai.ai.auto_ai
Expand source code
from __future__ import annotations
from typing import Optional
from ..messages import AIMessages, AIMessage
from gguf_modeldb import ModelDB, ModelData
from gguf_llama import LlamaAI
__all__ = ['AutoAI']
# CONFIG
class AutoAI:
"""
Initialize the AutoAI class for super easy LLM AI generation.
Searches for a model based on name/quantization/keyword.
Downloads model and sets up LlamaAI.
Args:
name_search: Name of model to search for. Optional. Default None.
quantization_search: Quantization of model to search for. Optional. Default None.
keyword_search: Keyword of model to search for. Optional. Default None.
new_tokens: New token length for LlamaAI model. Default 1500.
max_input_tokens: Max input tokens for LlamaAI model. Default 900.
model_db_dir: Directory to store model data in. Defaults to global packages model directory.
Attributes:
model_db: ModelDB object. - represents the database of models, has useful functions for searching and importing models.
model_data: ModelData object. - represents the data of the model, has useful functions for creating, downloading and loading the model data and gguf.
ai: LlamaAI object. - represents the LlamaAI model, a wrapper for llama llm and tokenizer models quantized to gguf format. Has methods for adjusting generation and for generating.
msgs: AIMessages object. - represents the AIMessages a collection of AIMessage objects, has useful functions for adding and editing messages and can be printed to string.
"""
def __init__(self,
name_search: Optional[str] = None,
quantization_search: Optional[str] = None,
keyword_search: Optional[str] = None,
search_only_downloaded_models:bool = False,
max_total_tokens: int = 1500,
model_db_dir:Optional[str] = None,
) -> None:
self.model_db = ModelDB(model_db_dir=model_db_dir, copy_verified_models=True)
self.model_data: ModelData = self.model_db.find_model(
name_search, quantization_search, keyword_search, search_only_downloaded_models
)
self.model_data.download_gguf()
self.ai = LlamaAI(
self.model_data.gguf_file_path, max_tokens=max_total_tokens
)
print(f"Using model: {self.model_data}")
self.msgs: AIMessages = AIMessages(
self.model_data.user_tags, self.model_data.ai_tags, self.model_data.system_tags
)
def generate_from_messages(self, stop_at:str = None, include_stop_str:bool = True) -> AIMessage:
prompt = self.msgs.text()
ai_message = self.generate_from_literal_string(prompt, stop_at=stop_at, include_stop_str=include_stop_str)
self.msgs.add_ai_message(ai_message)
return ai_message
def generate_from_literal_string(
self,
prompt: str,
stop_at:str = None,
include_stop_str:bool = True
) -> AIMessage:
"""
Generate text from a prompt using the LlamaAI model.
Args:
prompt: Prompt text to generate from.
Returns:
Generated text string.
"""
return self.ai.infer(
prompt, only_string=True, stop_at_str=stop_at, include_stop_str=include_stop_str
)
def generate(
self,
user_message: str,
ai_message_tbc: Optional[str] = None,
stop_at:Optional[str] = None,
include_stop_str:bool = True,
system_message: Optional[str] = None
) -> AIMessage:
"""
Generate an AI response to a user message.
Args:
user_message: User message text.
ai_message_tbc: Optional text to prepend.
stop_at: Optional string to stop generation at.
include_stop_str: Whether to include the stop string in the generated message.
system_message: Optional system message to include at the start, not all models support this.
If you provide system message to a model that doesn't support it, it will be ignored.
You can check if a model supports system messages by checking the model_data.has_system_tags() method.
Returns:
Generated AIMessage object.
"""
generation_messages = AIMessages(user_tags=self.model_data.user_tags, ai_tags=self.model_data.ai_tags, system_tags=self.model_data.system_tags)
generation_messages.reset_messages()
if self.model_data.has_system_tags():
if system_message is not None:
generation_messages.set_system_message(system_message)
else:
print("WARNING: Model seeps to support system messages, but no system message provided.")
generation_messages.add_user_message(user_message)
if ai_message_tbc is not None:
generation_messages.add_message(
ai_message_tbc,
self.msgs.ai_tag_open,
""
)
print(f"Promt: {generation_messages.text()}")
generated = self.generate_from_literal_string(generation_messages.text(), stop_at=stop_at, include_stop_str=include_stop_str)
if ai_message_tbc is not None:
generation_messages.edit_last_message(
ai_message_tbc + generated,
self.msgs.ai_tag_open,
self.msgs.ai_tag_close
)
else:
generation_messages.add_ai_message(generated)
print(f"Generated: {generation_messages.get_last_message().text()}")
output = generation_messages.get_last_message()
return output
def count_tokens(
self,
user_message: str,
ai_message_tbc: Optional[str] = None
) -> int:
"""
Count the number of tokens in a generated message.
Args:
user_message: User message text.
ai_message_tbc: Optional text to prepend.
Returns:
Number of tokens in generated message.
"""
generation_messages = AIMessages()
generation_messages.reset_messages()
generation_messages.add_user_message(user_message)
if ai_message_tbc is not None:
generation_messages.add_message(
ai_message_tbc,
self.msgs.ai_tag_open,
""
)
return self.ai.count_tokens(generation_messages.text())
def is_within_input_limit(
self,
user_message: str,
ai_message_tbc: Optional[str] = None
) -> bool:
"""
Check if the generated message is within the input limit.
Args:
user_message: User message text.
ai_message_tbc: Optional text to prepend.
Returns:
True if within input limit, False otherwise.
"""
return self.ai.is_within_input_limit(self.count_tokens(user_message, ai_message_tbc))
Classes
class AutoAI (name_search: Optional[str] = None, quantization_search: Optional[str] = None, keyword_search: Optional[str] = None, search_only_downloaded_models: bool = False, max_total_tokens: int = 1500, model_db_dir: Optional[str] = None)
-
Initialize the AutoAI class for super easy LLM AI generation.
Searches for a model based on name/quantization/keyword. Downloads model and sets up LlamaAI.
Args
name_search
- Name of model to search for. Optional. Default None.
quantization_search
- Quantization of model to search for. Optional. Default None.
keyword_search
- Keyword of model to search for. Optional. Default None.
new_tokens
- New token length for LlamaAI model. Default 1500.
max_input_tokens
- Max input tokens for LlamaAI model. Default 900.
model_db_dir
- Directory to store model data in. Defaults to global packages model directory.
Attributes
model_db
- ModelDB object. - represents the database of models, has useful functions for searching and importing models.
model_data
- ModelData object. - represents the data of the model, has useful functions for creating, downloading and loading the model data and gguf.
ai
- LlamaAI object. - represents the LlamaAI model, a wrapper for llama llm and tokenizer models quantized to gguf format. Has methods for adjusting generation and for generating.
msgs
- AIMessages object. - represents the AIMessages a collection of AIMessage objects, has useful functions for adding and editing messages and can be printed to string.
Expand source code
class AutoAI: """ Initialize the AutoAI class for super easy LLM AI generation. Searches for a model based on name/quantization/keyword. Downloads model and sets up LlamaAI. Args: name_search: Name of model to search for. Optional. Default None. quantization_search: Quantization of model to search for. Optional. Default None. keyword_search: Keyword of model to search for. Optional. Default None. new_tokens: New token length for LlamaAI model. Default 1500. max_input_tokens: Max input tokens for LlamaAI model. Default 900. model_db_dir: Directory to store model data in. Defaults to global packages model directory. Attributes: model_db: ModelDB object. - represents the database of models, has useful functions for searching and importing models. model_data: ModelData object. - represents the data of the model, has useful functions for creating, downloading and loading the model data and gguf. ai: LlamaAI object. - represents the LlamaAI model, a wrapper for llama llm and tokenizer models quantized to gguf format. Has methods for adjusting generation and for generating. msgs: AIMessages object. - represents the AIMessages a collection of AIMessage objects, has useful functions for adding and editing messages and can be printed to string. """ def __init__(self, name_search: Optional[str] = None, quantization_search: Optional[str] = None, keyword_search: Optional[str] = None, search_only_downloaded_models:bool = False, max_total_tokens: int = 1500, model_db_dir:Optional[str] = None, ) -> None: self.model_db = ModelDB(model_db_dir=model_db_dir, copy_verified_models=True) self.model_data: ModelData = self.model_db.find_model( name_search, quantization_search, keyword_search, search_only_downloaded_models ) self.model_data.download_gguf() self.ai = LlamaAI( self.model_data.gguf_file_path, max_tokens=max_total_tokens ) print(f"Using model: {self.model_data}") self.msgs: AIMessages = AIMessages( self.model_data.user_tags, self.model_data.ai_tags, self.model_data.system_tags ) def generate_from_messages(self, stop_at:str = None, include_stop_str:bool = True) -> AIMessage: prompt = self.msgs.text() ai_message = self.generate_from_literal_string(prompt, stop_at=stop_at, include_stop_str=include_stop_str) self.msgs.add_ai_message(ai_message) return ai_message def generate_from_literal_string( self, prompt: str, stop_at:str = None, include_stop_str:bool = True ) -> AIMessage: """ Generate text from a prompt using the LlamaAI model. Args: prompt: Prompt text to generate from. Returns: Generated text string. """ return self.ai.infer( prompt, only_string=True, stop_at_str=stop_at, include_stop_str=include_stop_str ) def generate( self, user_message: str, ai_message_tbc: Optional[str] = None, stop_at:Optional[str] = None, include_stop_str:bool = True, system_message: Optional[str] = None ) -> AIMessage: """ Generate an AI response to a user message. Args: user_message: User message text. ai_message_tbc: Optional text to prepend. stop_at: Optional string to stop generation at. include_stop_str: Whether to include the stop string in the generated message. system_message: Optional system message to include at the start, not all models support this. If you provide system message to a model that doesn't support it, it will be ignored. You can check if a model supports system messages by checking the model_data.has_system_tags() method. Returns: Generated AIMessage object. """ generation_messages = AIMessages(user_tags=self.model_data.user_tags, ai_tags=self.model_data.ai_tags, system_tags=self.model_data.system_tags) generation_messages.reset_messages() if self.model_data.has_system_tags(): if system_message is not None: generation_messages.set_system_message(system_message) else: print("WARNING: Model seeps to support system messages, but no system message provided.") generation_messages.add_user_message(user_message) if ai_message_tbc is not None: generation_messages.add_message( ai_message_tbc, self.msgs.ai_tag_open, "" ) print(f"Promt: {generation_messages.text()}") generated = self.generate_from_literal_string(generation_messages.text(), stop_at=stop_at, include_stop_str=include_stop_str) if ai_message_tbc is not None: generation_messages.edit_last_message( ai_message_tbc + generated, self.msgs.ai_tag_open, self.msgs.ai_tag_close ) else: generation_messages.add_ai_message(generated) print(f"Generated: {generation_messages.get_last_message().text()}") output = generation_messages.get_last_message() return output def count_tokens( self, user_message: str, ai_message_tbc: Optional[str] = None ) -> int: """ Count the number of tokens in a generated message. Args: user_message: User message text. ai_message_tbc: Optional text to prepend. Returns: Number of tokens in generated message. """ generation_messages = AIMessages() generation_messages.reset_messages() generation_messages.add_user_message(user_message) if ai_message_tbc is not None: generation_messages.add_message( ai_message_tbc, self.msgs.ai_tag_open, "" ) return self.ai.count_tokens(generation_messages.text()) def is_within_input_limit( self, user_message: str, ai_message_tbc: Optional[str] = None ) -> bool: """ Check if the generated message is within the input limit. Args: user_message: User message text. ai_message_tbc: Optional text to prepend. Returns: True if within input limit, False otherwise. """ return self.ai.is_within_input_limit(self.count_tokens(user_message, ai_message_tbc))
Methods
def count_tokens(self, user_message: str, ai_message_tbc: Optional[str] = None) ‑> int
-
Count the number of tokens in a generated message.
Args
user_message
- User message text.
ai_message_tbc
- Optional text to prepend.
Returns
Number of tokens in generated message.
Expand source code
def count_tokens( self, user_message: str, ai_message_tbc: Optional[str] = None ) -> int: """ Count the number of tokens in a generated message. Args: user_message: User message text. ai_message_tbc: Optional text to prepend. Returns: Number of tokens in generated message. """ generation_messages = AIMessages() generation_messages.reset_messages() generation_messages.add_user_message(user_message) if ai_message_tbc is not None: generation_messages.add_message( ai_message_tbc, self.msgs.ai_tag_open, "" ) return self.ai.count_tokens(generation_messages.text())
def generate(self, user_message: str, ai_message_tbc: Optional[str] = None, stop_at: Optional[str] = None, include_stop_str: bool = True, system_message: Optional[str] = None) ‑> AIMessage
-
Generate an AI response to a user message.
Args
user_message
- User message text.
ai_message_tbc
- Optional text to prepend.
stop_at
- Optional string to stop generation at.
include_stop_str
- Whether to include the stop string in the generated message.
system_message
- Optional system message to include at the start, not all models support this.
If you provide system message to a model that doesn't support it, it will be ignored. You can check if a model supports system messages by checking the model_data.has_system_tags() method.
Returns
Generated AIMessage object.
Expand source code
def generate( self, user_message: str, ai_message_tbc: Optional[str] = None, stop_at:Optional[str] = None, include_stop_str:bool = True, system_message: Optional[str] = None ) -> AIMessage: """ Generate an AI response to a user message. Args: user_message: User message text. ai_message_tbc: Optional text to prepend. stop_at: Optional string to stop generation at. include_stop_str: Whether to include the stop string in the generated message. system_message: Optional system message to include at the start, not all models support this. If you provide system message to a model that doesn't support it, it will be ignored. You can check if a model supports system messages by checking the model_data.has_system_tags() method. Returns: Generated AIMessage object. """ generation_messages = AIMessages(user_tags=self.model_data.user_tags, ai_tags=self.model_data.ai_tags, system_tags=self.model_data.system_tags) generation_messages.reset_messages() if self.model_data.has_system_tags(): if system_message is not None: generation_messages.set_system_message(system_message) else: print("WARNING: Model seeps to support system messages, but no system message provided.") generation_messages.add_user_message(user_message) if ai_message_tbc is not None: generation_messages.add_message( ai_message_tbc, self.msgs.ai_tag_open, "" ) print(f"Promt: {generation_messages.text()}") generated = self.generate_from_literal_string(generation_messages.text(), stop_at=stop_at, include_stop_str=include_stop_str) if ai_message_tbc is not None: generation_messages.edit_last_message( ai_message_tbc + generated, self.msgs.ai_tag_open, self.msgs.ai_tag_close ) else: generation_messages.add_ai_message(generated) print(f"Generated: {generation_messages.get_last_message().text()}") output = generation_messages.get_last_message() return output
def generate_from_literal_string(self, prompt: str, stop_at: str = None, include_stop_str: bool = True) ‑> AIMessage
-
Generate text from a prompt using the LlamaAI model.
Args
prompt
- Prompt text to generate from.
Returns
Generated text string.
Expand source code
def generate_from_literal_string( self, prompt: str, stop_at:str = None, include_stop_str:bool = True ) -> AIMessage: """ Generate text from a prompt using the LlamaAI model. Args: prompt: Prompt text to generate from. Returns: Generated text string. """ return self.ai.infer( prompt, only_string=True, stop_at_str=stop_at, include_stop_str=include_stop_str )
def generate_from_messages(self, stop_at: str = None, include_stop_str: bool = True) ‑> AIMessage
-
Expand source code
def generate_from_messages(self, stop_at:str = None, include_stop_str:bool = True) -> AIMessage: prompt = self.msgs.text() ai_message = self.generate_from_literal_string(prompt, stop_at=stop_at, include_stop_str=include_stop_str) self.msgs.add_ai_message(ai_message) return ai_message
def is_within_input_limit(self, user_message: str, ai_message_tbc: Optional[str] = None) ‑> bool
-
Check if the generated message is within the input limit.
Args
user_message
- User message text.
ai_message_tbc
- Optional text to prepend.
Returns
True if within input limit, False otherwise.
Expand source code
def is_within_input_limit( self, user_message: str, ai_message_tbc: Optional[str] = None ) -> bool: """ Check if the generated message is within the input limit. Args: user_message: User message text. ai_message_tbc: Optional text to prepend. Returns: True if within input limit, False otherwise. """ return self.ai.is_within_input_limit(self.count_tokens(user_message, ai_message_tbc))