aux_backup / tiny_styler.py
harvesthealth's picture
Upload folder using huggingface_hub
71d0e7d verified
from tinytroupe.enrichment import logger
from tinytroupe.utils import JsonSerializableRegistry
from tinytroupe.utils.llm import LLMChat
import tinytroupe.utils as utils
class TinyStyler(JsonSerializableRegistry):
"""
A class for applying a specified writing or speaking style to content while preserving
the original information.
"""
def __init__(self, use_past_results_in_context=False) -> None:
"""
Initialize the TinyStyler.
Args:
use_past_results_in_context (bool): Whether to use past styling results in the context.
"""
self.use_past_results_in_context = use_past_results_in_context
self.context_cache = []
def apply_style(self, content: str, style: str, content_type: str = None,
context_info: str = "", context_cache: list = None, verbose: bool = False,
temperature: float = 0.7):
"""
Apply a specified style to the content while preserving all the original information.
Args:
content (str): The content to style.
style (str): The style to apply (e.g., "professional", "casual", "technical", etc.).
content_type (str, optional): The type of content (e.g., "email", "report", "conversation").
context_info (str, optional): Additional context information.
context_cache (list, optional): Previous styling results to use as context.
verbose (bool, optional): Whether to print debug information.
temperature (float, optional): The temperature to use for the LLM generation.
Returns:
str: The styled content.
"""
if context_cache is None and self.use_past_results_in_context:
context_cache = self.context_cache
rendering_configs = {
"content": content,
"style": style,
"content_type": content_type,
"context_info": context_info,
"context_cache": context_cache
}
# Initialize the LLMChat with appropriate templates
chat = LLMChat(
system_template_name="styler.system.mustache",
user_template_name="styler.user.mustache",
base_module_folder="enrichment",
temperature=temperature
)
# Call the model and get the response
result = chat.call(**rendering_configs)
debug_msg = f"Styling result: {result}"
logger.debug(debug_msg)
if verbose:
print(debug_msg)
# Extract the styled content from code blocks if present
if result is not None:
styled_content = utils.extract_code_block(result)
# If no code block was found, use the raw result
if not styled_content:
styled_content = result
# Add to context cache if enabled
if self.use_past_results_in_context:
self.context_cache.append({
"original": content,
"style": style,
"styled": styled_content
})
return styled_content
else:
return None