PraisonAI 0.0.59rc2__cp312-cp312-manylinux_2_35_x86_64.whl

praisonai/train.py

@@ -0,0 +1,232 @@
+import subprocess
+import os
+import sys
+import yaml
+import torch
+from transformers import TextStreamer
+from unsloth import FastLanguageModel, is_bfloat16_supported
+from trl import SFTTrainer
+from transformers import TrainingArguments
+from datasets import load_dataset, concatenate_datasets, Dataset
+from psutil import virtual_memory
+
+class train:
+    def __init__(self, config_path="config.yaml"):
+        self.load_config(config_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model, self.tokenizer = None, None
+
+    def load_config(self, path):
+        with open(path, "r") as file:
+            self.config = yaml.safe_load(file)
+
+    def print_system_info(self):
+        print(f"PyTorch version: {torch.__version__}")
+        print(f"CUDA version: {torch.version.cuda}")
+        if torch.cuda.is_available():
+            device_capability = torch.cuda.get_device_capability()
+            print(f"CUDA Device Capability: {device_capability}")
+        else:
+            print("CUDA is not available")
+
+        python_version = sys.version
+        pip_version = subprocess.check_output(['pip', '--version']).decode().strip()
+        python_path = sys.executable
+        pip_path = subprocess.check_output(['which', 'pip']).decode().strip()
+        print(f"Python Version: {python_version}")
+        print(f"Pip Version: {pip_version}")
+        print(f"Python Path: {python_path}")
+        print(f"Pip Path: {pip_path}")
+
+    def check_gpu(self):
+        gpu_stats = torch.cuda.get_device_properties(0)
+        print(f"GPU = {gpu_stats.name}. Max memory = {round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)} GB.")
+
+    def check_ram(self):
+        ram_gb = virtual_memory().total / 1e9
+        print('Your runtime has {:.1f} gigabytes of available RAM\n'.format(ram_gb))
+        if ram_gb < 20:
+            print('Not using a high-RAM runtime')
+        else:
+            print('You are using a high-RAM runtime!')
+
+    # def install_packages(self):
+    #     subprocess.run(["pip", "install", "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git@4e570be9ae4ced8cdc64e498125708e34942befc"])
+    #     subprocess.run(["pip", "install", "--no-deps", "trl<0.9.0", "peft==0.12.0", "accelerate==0.33.0", "bitsandbytes==0.43.3"])
+
+    def prepare_model(self):
+        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
+            model_name=self.config["model_name"],
+            max_seq_length=self.config["max_seq_length"],
+            dtype=None,
+            load_in_4bit=self.config["load_in_4bit"]
+        )
+        self.model = FastLanguageModel.get_peft_model(
+            self.model,
+            r=self.config["lora_r"],
+            target_modules=self.config["lora_target_modules"],
+            lora_alpha=self.config["lora_alpha"],
+            lora_dropout=self.config["lora_dropout"],
+            bias=self.config["lora_bias"],
+            use_gradient_checkpointing=self.config["use_gradient_checkpointing"],
+            random_state=self.config["random_state"],
+            use_rslora=self.config["use_rslora"],
+            loftq_config=self.config["loftq_config"],
+        )
+
+    def process_dataset(self, dataset_info):
+        dataset_name = dataset_info["name"]
+        split_type = dataset_info.get("split_type", "train")
+        processing_func = getattr(self, dataset_info.get("processing_func", "format_prompts"))
+        rename = dataset_info.get("rename", {})
+        filter_data = dataset_info.get("filter_data", False)
+        filter_column_value = dataset_info.get("filter_column_value", "id")
+        filter_value = dataset_info.get("filter_value", "alpaca")
+        num_samples = dataset_info.get("num_samples", 20000)
+
+        dataset = load_dataset(dataset_name, split=split_type)
+
+        if rename:
+            dataset = dataset.rename_columns(rename)
+        if filter_data:
+            dataset = dataset.filter(lambda example: filter_value in example[filter_column_value]).shuffle(seed=42).select(range(num_samples))
+        dataset = dataset.map(processing_func, batched=True)
+        return dataset
+
+    def format_prompts(self, examples):
+        alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+        ### Instruction:
+        {}
+
+        ### Input:
+        {}
+
+        ### Response:
+        {}"""
+        texts = [alpaca_prompt.format(ins, inp, out) + self.tokenizer.eos_token for ins, inp, out in zip(examples["instruction"], examples["input"], examples["output"])]
+        return {"text": texts}
+
+    def load_datasets(self):
+        datasets = []
+        for dataset_info in self.config["dataset"]:
+            datasets.append(self.process_dataset(dataset_info))
+        return concatenate_datasets(datasets)
+
+    def train_model(self):
+        dataset = self.load_datasets()
+        trainer = SFTTrainer(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            train_dataset=dataset,
+            dataset_text_field=self.config["dataset_text_field"],
+            max_seq_length=self.config["max_seq_length"],
+            dataset_num_proc=self.config["dataset_num_proc"],
+            packing=self.config["packing"],
+            args=TrainingArguments(
+                per_device_train_batch_size=self.config["per_device_train_batch_size"],
+                gradient_accumulation_steps=self.config["gradient_accumulation_steps"],
+                warmup_steps=self.config["warmup_steps"],
+                num_train_epochs=self.config["num_train_epochs"],
+                max_steps=self.config["max_steps"],
+                learning_rate=self.config["learning_rate"],
+                fp16=not is_bfloat16_supported(),
+                bf16=is_bfloat16_supported(),
+                logging_steps=self.config["logging_steps"],
+                optim=self.config["optim"],
+                weight_decay=self.config["weight_decay"],
+                lr_scheduler_type=self.config["lr_scheduler_type"],
+                seed=self.config["seed"],
+                output_dir=self.config["output_dir"],
+            ),
+        )
+        trainer.train()
+
+    def inference(self, instruction, input_text):
+        FastLanguageModel.for_inference(self.model)
+        alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+        ### Instruction:
+        {}
+
+        ### Input:
+        {}
+
+        ### Response:
+        {}"""
+        inputs = self.tokenizer([alpaca_prompt.format(instruction, input_text, "")], return_tensors="pt").to("cuda")
+        outputs = self.model.generate(**inputs, max_new_tokens=64, use_cache=True)
+        print(self.tokenizer.batch_decode(outputs))
+
+    def save_model_merged(self):
+        if os.path.exists(self.config["hf_model_name"]):
+            shutil.rmtree(self.config["hf_model_name"])
+        self.model.push_to_hub_merged(
+            self.config["hf_model_name"],
+            self.tokenizer,
+            save_method="merged_16bit",
+            token=os.getenv('HF_TOKEN')
+        )
+
+    def push_model_gguf(self):
+        self.model.push_to_hub_gguf(
+            self.config["hf_model_name"],
+            self.tokenizer,
+            quantization_method=self.config["quantization_method"],
+            token=os.getenv('HF_TOKEN')
+        )
+
+    def prepare_modelfile_content(self):
+        output_model = self.config["hf_model_name"]
+        return f"""FROM {output_model}/unsloth.Q5_K_M.gguf
+
+TEMPLATE \"\"\"Below are some instructions that describe some tasks. Write responses that appropriately complete each request.{{{{ if .Prompt }}}}
+
+### Instruction:
+{{{{ .Prompt }}}}
+
+{{{{ end }}}}### Response:
+{{{{ .Response }}}}\"\"\"
+
+PARAMETER stop ""
+PARAMETER stop ""
+PARAMETER stop ""
+PARAMETER stop ""
+PARAMETER stop "<|reserved_special_token_"
+"""
+
+    def create_and_push_ollama_model(self):
+        modelfile_content = self.prepare_modelfile_content()
+        with open('Modelfile', 'w') as file:
+            file.write(modelfile_content)
+
+        subprocess.run(["ollama", "serve"])
+        subprocess.run(["ollama", "create", f"{self.config['ollama_model']}:{self.config['model_parameters']}", "-f", "Modelfile"])
+        subprocess.run(["ollama", "push", f"{self.config['ollama_model']}:{self.config['model_parameters']}"])
+
+    def run(self):
+        self.print_system_info()
+        self.check_gpu()
+        self.check_ram()
+        # self.install_packages()
+        self.prepare_model()
+        self.train_model()
+        self.save_model_merged()
+        self.push_model_gguf()
+        self.create_and_push_ollama_model()
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description='PraisonAI Training Script')
+    parser.add_argument('command', choices=['train'], help='Command to execute')
+    parser.add_argument('--config', default='config.yaml', help='Path to configuration file')
+    args = parser.parse_args()
+
+    if args.command == 'train':
+        ai = train(config_path=args.config)
+        ai.run()
+
+
+if __name__ == '__main__':
+    main()

praisonai/ui/chat.py

@@ -0,0 +1,304 @@
+import chainlit as cl
+from chainlit.input_widget import TextInput
+from chainlit.types import ThreadDict
+from litellm import acompletion
+import os
+import sqlite3
+from datetime import datetime
+from typing import Dict, List, Optional
+from dotenv import load_dotenv
+load_dotenv()
+import chainlit.data as cl_data
+from chainlit.step import StepDict
+from literalai.helper import utc_now
+import logging
+import json
+from sql_alchemy import SQLAlchemyDataLayer
+
+# Set up logging
+logger = logging.getLogger(__name__)
+log_level = os.getenv("LOGLEVEL", "INFO").upper()
+logger.handlers = []
+
+# Set up logging to console
+console_handler = logging.StreamHandler()
+console_handler.setLevel(log_level)
+console_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+console_handler.setFormatter(console_formatter)
+logger.addHandler(console_handler)
+
+# Set the logging level for the logger
+logger.setLevel(log_level)
+
+CHAINLIT_AUTH_SECRET = os.getenv("CHAINLIT_AUTH_SECRET")
+
+if not CHAINLIT_AUTH_SECRET:
+    os.environ["CHAINLIT_AUTH_SECRET"] = "p8BPhQChpg@J>jBz$wGxqLX2V>yTVgP*7Ky9H$aV:axW~ANNX-7_T:o@lnyCBu^U"
+    CHAINLIT_AUTH_SECRET = os.getenv("CHAINLIT_AUTH_SECRET")
+
+now = utc_now()
+
+create_step_counter = 0
+
+DB_PATH = os.path.expanduser("~/.praison/database.sqlite")
+
+def initialize_db():
+    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS users (
+            id UUID PRIMARY KEY,
+            identifier TEXT NOT NULL UNIQUE,
+            metadata JSONB NOT NULL,
+            createdAt TEXT
+        )
+    ''')
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS threads (
+            id UUID PRIMARY KEY,
+            createdAt TEXT,
+            name TEXT,
+            userId UUID,
+            userIdentifier TEXT,
+            tags TEXT[],
+            metadata JSONB NOT NULL DEFAULT '{}',
+            FOREIGN KEY (userId) REFERENCES users(id) ON DELETE CASCADE
+        )
+    ''')
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS steps (
+            id UUID PRIMARY KEY,
+            name TEXT NOT NULL,
+            type TEXT NOT NULL,
+            threadId UUID NOT NULL,
+            parentId UUID,
+            disableFeedback BOOLEAN NOT NULL,
+            streaming BOOLEAN NOT NULL,
+            waitForAnswer BOOLEAN,
+            isError BOOLEAN,
+            metadata JSONB,
+            tags TEXT[],
+            input TEXT,
+            output TEXT,
+            createdAt TEXT,
+            start TEXT,
+            end TEXT,
+            generation JSONB,
+            showInput TEXT,
+            language TEXT,
+            indent INT,
+            FOREIGN KEY (threadId) REFERENCES threads (id) ON DELETE CASCADE
+        )
+    ''')
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS elements (
+            id UUID PRIMARY KEY,
+            threadId UUID,
+            type TEXT,
+            url TEXT,
+            chainlitKey TEXT,
+            name TEXT NOT NULL,
+            display TEXT,
+            objectKey TEXT,
+            size TEXT,
+            page INT,
+            language TEXT,
+            forId UUID,
+            mime TEXT,
+            FOREIGN KEY (threadId) REFERENCES threads (id) ON DELETE CASCADE
+        )
+    ''')
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS feedbacks (
+            id UUID PRIMARY KEY,
+            forId UUID NOT NULL,
+            value INT NOT NULL,
+            threadId UUID,
+            comment TEXT
+        )
+    ''')
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS settings (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            key TEXT UNIQUE,
+            value TEXT
+        )
+    ''')
+    conn.commit()
+    conn.close()
+
+def save_setting(key: str, value: str):
+    """Saves a setting to the database.
+
+    Args:
+        key: The setting key.
+        value: The setting value.
+    """
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute(
+        """
+        INSERT OR REPLACE INTO settings (id, key, value)
+        VALUES ((SELECT id FROM settings WHERE key = ?), ?, ?)
+    """,
+        (key, key, value),
+    )
+    conn.commit()
+    conn.close()
+
+def load_setting(key: str) -> str:
+    """Loads a setting from the database.
+
+    Args:
+        key: The setting key.
+
+    Returns:
+        The setting value, or None if the key is not found.
+    """
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute('SELECT value FROM settings WHERE key = ?', (key,))
+    result = cursor.fetchone()
+    conn.close()
+    return result[0] if result else None
+
+
+# Initialize the database
+initialize_db()
+
+deleted_thread_ids = []  # type: List[str]
+
+cl_data._data_layer = SQLAlchemyDataLayer(conninfo=f"sqlite+aiosqlite:///{DB_PATH}")
+
+@cl.on_chat_start
+async def start():
+    initialize_db()
+    model_name = load_setting("model_name") 
+
+    if model_name:
+        cl.user_session.set("model_name", model_name)
+    else:
+        # If no setting found, use default or environment variable
+        model_name = os.getenv("MODEL_NAME", "gpt-4o-mini")
+        cl.user_session.set("model_name", model_name)
+    logger.debug(f"Model name: {model_name}")
+    settings = cl.ChatSettings(
+        [
+            TextInput(
+                id="model_name",
+                label="Enter the Model Name",
+                placeholder="e.g., gpt-4o-mini",
+                initial=model_name
+            )
+        ]
+    )
+    cl.user_session.set("settings", settings)
+    await settings.send()
+
+@cl.on_settings_update
+async def setup_agent(settings):
+    logger.debug(settings)
+    cl.user_session.set("settings", settings)
+    model_name = settings["model_name"]
+    cl.user_session.set("model_name", model_name)
+    
+    # Save in settings table
+    save_setting("model_name", model_name)
+    
+    # Save in thread metadata
+    thread_id = cl.user_session.get("thread_id")
+    if thread_id:
+        thread = await cl_data.get_thread(thread_id)
+        if thread:
+            metadata = thread.get("metadata", {})
+            metadata["model_name"] = model_name
+            
+            # Always store metadata as a JSON string
+            await cl_data.update_thread(thread_id, metadata=json.dumps(metadata))
+            
+            # Update the user session with the new metadata
+            cl.user_session.set("metadata", metadata)
+
+@cl.on_message
+async def main(message: cl.Message):
+    model_name = load_setting("model_name") or os.getenv("MODEL_NAME") or "gpt-4o-mini"
+    message_history = cl.user_session.get("message_history", [])
+    message_history.append({"role": "user", "content": message.content})
+
+    msg = cl.Message(content="")
+    await msg.send()
+
+    response = await acompletion(
+        model=model_name,
+        messages=message_history,
+        stream=True,
+        # temperature=0.7,
+        # max_tokens=500,
+        # top_p=1
+    )
+
+    full_response = ""
+    async for part in response:
+        if token := part['choices'][0]['delta']['content']:
+            await msg.stream_token(token)
+            full_response += token
+    logger.debug(f"Full response: {full_response}")
+    message_history.append({"role": "assistant", "content": full_response})
+    logger.debug(f"Message history: {message_history}")
+    cl.user_session.set("message_history", message_history)
+    await msg.update()
+
+username = os.getenv("CHAINLIT_USERNAME", "admin")  # Default to "admin" if not found
+password = os.getenv("CHAINLIT_PASSWORD", "admin")  # Default to "admin" if not found
+
+@cl.password_auth_callback
+def auth_callback(username: str, password: str):
+    if (username, password) == (username, password):
+        return cl.User(
+            identifier=username, metadata={"role": "ADMIN", "provider": "credentials"}
+        )
+    else:
+        return None
+
+async def send_count():
+    await cl.Message(
+        f"Create step counter: {create_step_counter}", disable_feedback=True
+    ).send()
+
+@cl.on_chat_resume
+async def on_chat_resume(thread: cl_data.ThreadDict):
+    logger.info(f"Resuming chat: {thread['id']}")
+    model_name = load_setting("model_name") or os.getenv("MODEL_NAME") or "gpt-4o-mini"
+    logger.debug(f"Model name: {model_name}")
+    settings = cl.ChatSettings(
+        [
+            TextInput(
+                id="model_name",
+                label="Enter the Model Name",
+                placeholder="e.g., gpt-4o-mini",
+                initial=model_name
+            )
+        ]
+    )
+    await settings.send()
+    thread_id = thread["id"]
+    cl.user_session.set("thread_id", thread["id"])
+    
+    # The metadata should now already be a dictionary
+    metadata = thread.get("metadata", {})
+    cl.user_session.set("metadata", metadata)
+    
+    message_history = cl.user_session.get("message_history", [])
+    steps = thread["steps"]
+
+    for message in steps:
+        msg_type = message.get("type")
+        if msg_type == "user_message":
+            message_history.append({"role": "user", "content": message.get("output", "")})
+        elif msg_type == "assistant_message":
+            message_history.append({"role": "assistant", "content": message.get("output", "")})
+        else:
+            logger.warning(f"Message without type: {message}")
+
+    cl.user_session.set("message_history", message_history)