feat(models): add model tuning support

dineshpinto · dineshpinto · commit d34cb37a089e · 2025-02-16T13:11:42.000+04:00
diff --git a/.env.example b/.env.example
@@ -1,2 +1,7 @@
 # Local variables
-GEMINI_API_KEY=your-api-key
+GEMINI_API_KEY=your-api-key
+# Tuning parameters
+TUNING_SOURCE_MODEL=models/gemini-1.5-flash-001-tuning
+TUNING_EPOCH_COUNT=100
+TUNING_BATCH_SIZE=4
+TUNING_LEARNING_RATE=0.001 
diff --git a/README.md b/README.md
@@ -1 +1,20 @@
-# Flare AI Social
+# Flare AI Social
+
+## 🏗️ Build & Run Instructions
+
+You can deploy Flare AI DeFAI using Docker (recommended) or set up the backend and frontend manually.
+
+1. **Prepare the Environment File:**  
+   Rename `.env.example` to `.env` and update the variables accordingly.
+
+2. **Tune a new model**
+
+   ```bash
+   uv run start-tuning
+   ```
+
+3. **Test the new model**
+
+   ```bash
+   uv run start-social
+   ```
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,6 +20,7 @@ dependencies = [
 
 [project.scripts]
 start-social = "flare_ai_social.main:start"
+start-tuning = "flare_ai_social.tune_model:start"
 
 [build-system]
 requires = ["hatchling"]
diff --git a/src/data/training_data.json b/src/data/training_data.json
@@ -0,0 +1,26 @@
+[
+  {
+    "text_input": ".@Ripple x @Chainlink: $RLUSD has adopted the Chainlink standard for verifiable data to fuel DeFi adoption with real-time, secure pricing data. The future of stablecoins is here",
+    "output": "I'm at a loss as to why this would be interesting. Ripple needs an oracle so that RLUSD can be used in lending protocols on Ethereum. Flare doesn't provide oracles to other networks yet. It's something that may come but there are much bigger opportunities to pursue first: FAssets and native staking for XRP and BTC to name but two."
+  },
+  {
+    "text_input": "Man I hope I didn't make a mistake in going heavy with #flare and the @HugoPhilion team. Looking like U.S. based #Crypto like #XRP will get tax breaks on capital gains. Hugo,  any plans to move to the USA?",
+    "output": "No one has any knowledge of the rules around this yet. Rest assured if there is an opportunity for us to make sure that our American community aren't penalised for holding Flare relative to other tokens we will take it."
+  },
+  {
+    "text_input": "Where can I short $TRUMP Coin? 😂",
+    "output": "On Flare pretty soon you will be able to go long or short with leverage via @SparkDexAI perps."
+  },
+  {
+    "text_input": "Uhhh, sorry guys, did we forget we are building the tech for the future? Some decentralized system that will benefit the people of this planet and save them from tyranical govs. banks and powerful multinationals? It wasn't meant to be used for thousands of memecoins... hello...",
+    "output": "I understand that it looks like the market is going in the wrong direction but zoom out. Billions of dollars of value are assigned to things that often seem frivolous, such as art works, jewellery and fashion. There is no reason to think that the same human impulses that give these items value will not also be at play in a decentralized setting. Flare exists to allow people to build what they wish in the best way possible with maximum decentralization, whether that is perps on a hot meme coin, institutional finance or the future of AI. We are here for all of it."
+  },
+  {
+    "text_input": "Regarding enshrined Oracles like Flare, Supra, et al... Conceptually I love it. It's a strong business assumption.  But as a dev there are definitely scenarios where I'd want features of other L1s that the enshrined L1s do not offer. Whomever lands first with an Irresistible L1, epic marketing, easy onboarding, and a lost cost, easy-to-use, trust-less Oracle system - may just end up own a gigantic eco-system. The potential is huge. I can see devs flocking here 📈",
+    "output": "What else do you want to see on Flare?"
+  },
+  {
+    "text_input": "For all my coffee snobs out there what is the absolute best coffee/espresso machine?",
+    "output": "Breville Touch if you want an easy life. Decent DE1Pro if you want to go mental."
+  }
+]
diff --git a/src/flare_ai_social/__init__.py b/src/flare_ai_social/__init__.py
@@ -1,2 +0,0 @@
-def main() -> None:
-    print("Hello from flare-ai-social!")
diff --git a/src/flare_ai_social/ai/base.py b/src/flare_ai_social/ai/base.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any, Literal, Protocol, TypedDict, runtime_checkable
+from typing import Any, Protocol, TypedDict, runtime_checkable
 
 import httpx
 import requests
@@ -74,19 +74,23 @@ def send_message(self, msg: str) -> ModelResponse:
         """
 
 
+class Message(TypedDict):
+    role: str
+    content: str
+
+
 class CompletionRequest(TypedDict):
     model: str
     prompt: str
-
-
-class Message(TypedDict):
-    role: Literal["user", "assistant", "system"]
-    content: str
+    max_tokens: int
+    temperature: float
 
 
 class ChatRequest(TypedDict):
     model: str
     messages: list[Message]
+    max_tokens: int
+    temperature: float
 
 
 class BaseRouter:
diff --git a/src/flare_ai_social/main.py b/src/flare_ai_social/main.py
@@ -7,9 +7,7 @@
 genai.configure(api_key=settings.gemini_api_key)
 
 
-def start(tuned_model_id: str) -> None:
-    tuned_model_id = "pugo-hillion"
-
+def start(tuned_model_id: str = "pugo-hillion") -> None:
     tuned_models = [m.name for m in genai.list_tuned_models()]
     logger.info("available tuned models", tuned_models=tuned_models)
 
@@ -27,4 +25,4 @@ def start(tuned_model_id: str) -> None:
 
 
 if __name__ == "__main__":
-    start(tuned_model_id="pugo-hillion")
+    start()
diff --git a/src/flare_ai_social/prompts/__init__.py b/src/flare_ai_social/prompts/__init__.py
@@ -1,4 +1,3 @@
-from .dataset import training_dataset
 from .templates import SYSTEM_INSTRUCTION
 
-__all__ = ["SYSTEM_INSTRUCTION", "training_dataset"]
+__all__ = ["SYSTEM_INSTRUCTION"]
diff --git a/src/flare_ai_social/prompts/dataset.py b/src/flare_ai_social/prompts/dataset.py
diff --git a/src/flare_ai_social/settings.py b/src/flare_ai_social/settings.py
@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import structlog
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
@@ -11,6 +13,18 @@ class Settings(BaseSettings):
 
     # API key for accessing Google's Gemini AI service
     gemini_api_key: str = ""
+    # Tuning dataset path
+    tuning_dataset_path: Path = (
+        Path(__file__).parent.parent / "data" / "training_data.json"
+    )
+    # Base model to tune upon
+    tuning_source_model: str = "models/gemini-1.5-flash-001-tuning"
+    # Number of epochs to tune for
+    tuning_epoch_count: int = 100
+    # Batch size
+    tuning_batch_size: int = 4
+    # Learning rate
+    tuning_learning_rate: float = 0.001
 
     model_config = SettingsConfigDict(
         env_file=".env",
diff --git a/src/flare_ai_social/tune_model.py b/src/flare_ai_social/tune_model.py
@@ -1,53 +1,154 @@
+"""
+Module for training and tuning AI models using Google's Generative AI API.
+Handles model creation, training, and performance visualization.
+"""
+
+import json
+from pathlib import Path
+from typing import TypedDict
+
 import google.generativeai as genai
+import matplotlib.pyplot as plt
 import pandas as pd
 import seaborn as sns
 import structlog
 
-from flare_ai_social.prompts import training_dataset
 from flare_ai_social.settings import settings
 
 logger = structlog.get_logger(__name__)
 genai.configure(api_key=settings.gemini_api_key)
 
 
-def list_models_supporting_tuning() -> None:
-    models = [m.name for m in genai.list_models() if "createTunedModel" in m.supported_generation_methods]
-    logger.info("models supporting tuning supported", models=models)
+class TrainingEntry(TypedDict):
+    text_input: str
+    output: str
+
+
+TrainingData = list[TrainingEntry]
+
+
+def get_tuning_supported_models() -> list[str]:
+    """
+    Retrieve list of models that support fine-tuning.
+
+    Returns:
+        list[str]: Names of models supporting tuning
+    """
+    models = [
+        m.name
+        for m in genai.list_models()
+        if "createTunedModel" in m.supported_generation_methods
+    ]
+    logger.info("models supporting tuning identified", models=models)
+    return models
 
 
-def start(new_model_id: str) -> None:
-    # Check if model exists and delete it
+def delete_existing_model(model_id: str) -> None:
+    """
+    Delete a tuned model if it exists.
+
+    Args:
+        model_id (str): ID of the model to delete
+    """
+    full_model_id = f"tunedModels/{model_id}"
     for tuned_model in genai.list_tuned_models():
-        if tuned_model.name == f"tunedModels/{new_model_id}":
-            logger.info("deleting existing model", tuned_model_id=new_model_id)
-            genai.delete_tuned_model(f"tunedModels/{new_model_id}")
+        if tuned_model.name == full_model_id:
+            logger.info("deleting existing model", tuned_model_id=model_id)
+            genai.delete_tuned_model(full_model_id)
+            break
+
+
+def load_training_data(path: Path) -> TrainingData:
+    """
+    Load training dataset from a JSON file.
+
+    Args:
+        path (Path): Path to the training dataset file
+
+    Returns:
+        TrainingData: Loaded training dataset
+
+    Raises:
+        JSONDecodeError: If JSON file is invalid
+        FileNotFoundError: If file doesn't exist
+    """
+    try:
+        with path.open() as f:
+            return json.load(f)
+    except (json.JSONDecodeError, FileNotFoundError) as e:
+        logger.Exception("failed to load training data", error=str(e))
+        raise
 
-    # Create tuning operation
+
+def save_loss_plot(
+    snapshots: pd.DataFrame, model_id: str, save_path: str | None = None
+) -> Path:
+    """
+    Create and save a plot of mean loss over epochs.
+
+    Args:
+        snapshots (pd.DataFrame): Training snapshots data
+        model_id (str): ID of the model
+        save_path (Optional[str]): Custom save path for the plot
+
+    Returns:
+        Path: Path where the plot was saved
+    """
+    fig, ax = plt.subplots()
+    sns.set_style("darkgrid")
+
+    sns.lineplot(data=snapshots, x="epoch", y="mean_loss", ax=ax)
+    ax.set_title(f"Training Loss - {model_id}")
+    ax.set_xlabel("Epoch")
+    ax.set_ylabel("Mean Loss")
+
+    save_path = save_path or f"{model_id}_mean_loss.png"
+    fig.savefig(save_path)
+    plt.close(fig)
+
+    return Path(save_path)
+
+
+def start(new_model_id: str = "pugo-hillion") -> None:
+    """
+    Train a new model with the specified ID.
+
+    Args:
+        new_model_id (str): ID for the new model
+
+    Raises:
+        Exception: If model training fails
+    """
+    # Delete existing model if present
+    delete_existing_model(new_model_id)
+
+    # Load and validate training data
+    training_dataset = load_training_data(settings.tuning_dataset_path)
+
+    # Create and train model
+    logger.info("starting model tuning", tuned_model_id=new_model_id)
     operation = genai.create_tuned_model(
         id=new_model_id,
-        source_model="models/gemini-1.5-flash-001-tuning",
+        source_model=settings.tuning_source_model,
         training_data=training_dataset,
-        epoch_count=50,
-        batch_size=4,
-        learning_rate=0.005,
+        epoch_count=settings.tuning_epoch_count,
+        batch_size=settings.tuning_batch_size,
+        learning_rate=settings.tuning_learning_rate,
     )
 
-    # Wait for tuning to finish
+    # Monitor training progress
     logger.info("tuning model (takes a few mins)", tuned_model_id=new_model_id)
-    for _status in operation.wait_bar():
+    for _ in operation.wait_bar():
         pass
+
     tuned_model = operation.result()
     logger.info("tuning complete", tuned_model=tuned_model)
 
-    # Save mean loss to figure
+    # Generate and save loss plot
     snapshots = pd.DataFrame(tuned_model.tuning_task.snapshots)
-    plot = sns.lineplot(data=snapshots, x="epoch", y="mean_loss")
-    plot.set_title(new_model_id)
-    save_fig_path = f"{new_model_id}_mean_loss.png"
-    plot.get_figure().savefig(save_fig_path)  # pyright: ignore [reportAttributeAccessIssue, reportOptionalMemberAccess]
-    logger.info("saved mean_loss", save_fig_path=save_fig_path)
+    plot_path = save_loss_plot(snapshots, new_model_id)
+    logger.info("saved mean_loss plot", save_fig_path=str(plot_path))
 
 
 if __name__ == "__main__":
-    list_models_supporting_tuning()
-    # start("pugo-hillion")
+    start()

Original file line number	Diff line number	Diff line change
`@@ -1,2 +0,0 @@`
`1`		`-def main() -> None:`
`2`		`- print("Hello from flare-ai-social!")`