251 lines
8.0 KiB
Python
251 lines
8.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Step 5: Assign colors to speakers based on their characteristics.
|
|
|
|
Input: Speaker files in "_speakers/" folder
|
|
Output: _colors.json with speaker-color mappings
|
|
|
|
Output format:
|
|
{
|
|
"Malabar": "golden",
|
|
"Moon": "silver",
|
|
"Earth": "green",
|
|
...
|
|
}
|
|
|
|
Usage:
|
|
uv run step5_assign_colors.py
|
|
|
|
Environment Variables:
|
|
OPENAI_API_KEY - Required
|
|
OPENAI_BASE_URL - Optional (for Kimi/GLM APIs)
|
|
LLM_MODEL - Optional (e.g., "glm-4.5-air")
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
import sys
|
|
import json
|
|
from pathlib import Path
|
|
from typing import List, Dict, Tuple, Optional, Set
|
|
from openai import OpenAI
|
|
|
|
# ============== Configuration ==============
|
|
|
|
INPUT_DIR = Path("_speakers")
|
|
OUTPUT_FILE = Path("_colors.json")
|
|
|
|
# Fixed color assignments
|
|
FIXED_COLORS = {
|
|
"Malabar": "#FFD700" # Gold
|
|
}
|
|
|
|
# Default configurations for different providers
|
|
DEFAULT_CONFIGS = {
|
|
"openai": {
|
|
"base_url": None,
|
|
"model": "gpt-4o-mini"
|
|
},
|
|
"moonshot": {
|
|
"base_url": "https://api.moonshot.cn/v1",
|
|
"model": "kimi-latest"
|
|
},
|
|
"bigmodel": { # Zhipu AI (GLM)
|
|
"base_url": "https://open.bigmodel.cn/api/paas/v4",
|
|
"model": "glm-4.5-air"
|
|
}
|
|
}
|
|
|
|
|
|
def get_llm_config() -> Tuple[str, str]:
|
|
"""Get LLM configuration from environment."""
|
|
api_key = os.getenv("OPENAI_API_KEY")
|
|
if not api_key:
|
|
raise ValueError("OPENAI_API_KEY environment variable is required")
|
|
|
|
base_url = os.getenv("OPENAI_BASE_URL")
|
|
model = os.getenv("LLM_MODEL")
|
|
|
|
if base_url:
|
|
if model:
|
|
return base_url, model
|
|
if "bigmodel" in base_url:
|
|
return base_url, DEFAULT_CONFIGS["bigmodel"]["model"]
|
|
elif "moonshot" in base_url or "kimi" in base_url:
|
|
return base_url, DEFAULT_CONFIGS["moonshot"]["model"]
|
|
else:
|
|
return base_url, DEFAULT_CONFIGS["openai"]["model"]
|
|
else:
|
|
return None, model or DEFAULT_CONFIGS["openai"]["model"]
|
|
|
|
|
|
def collect_speakers(input_dir: Path) -> Set[str]:
|
|
"""Collect all unique speakers from speaker files."""
|
|
speakers = set()
|
|
|
|
for file_path in input_dir.glob("*_speakers.txt"):
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
# Parse line: [timestamp](Speaker) text
|
|
match = re.match(r'^\[\d{2}:\d{2}\]\(([^)]+)\)', line)
|
|
if match:
|
|
speakers.add(match.group(1))
|
|
|
|
return speakers
|
|
|
|
|
|
def assign_colors(speakers: Set[str], client: OpenAI, model: str) -> Dict[str, str]:
|
|
"""Assign colors to speakers using LLM."""
|
|
# Start with fixed colors
|
|
color_mapping = FIXED_COLORS.copy()
|
|
|
|
# Filter out speakers that already have fixed colors
|
|
remaining_speakers = [s for s in speakers if s not in color_mapping]
|
|
|
|
if not remaining_speakers:
|
|
return color_mapping
|
|
|
|
# Build prompt
|
|
speakers_list = ", ".join(remaining_speakers)
|
|
|
|
prompt = f"""Assign CSS hex color codes to each speaker from "Little Malabar" based on their characteristics.
|
|
|
|
Speakers to assign colors:
|
|
{speakers_list}
|
|
|
|
Color assignment guidelines (use hex codes like #FF0000):
|
|
- Mars → #CD5C5C (red planet) or #FF4500
|
|
- Earth → #228B22 (forest green) or #4169E1 (royal blue)
|
|
- Moon → #C0C0C0 (silver) or #A9A9A9 (dark gray)
|
|
- Sun → #FFD700 (gold) or #FFA500 (orange)
|
|
- Jupiter → #D2691E (chocolate/orange)
|
|
- Galaxy → #9370DB (medium purple) or #FF69B4 (hot pink)
|
|
- Star → #FFFFFF (white) or #FFFACD (lemon chiffon)
|
|
- Volcano → #8B0000 (dark red) or #FF4500 (orange red)
|
|
- Kangaroo/Giraffe → #D2B48C (tan) or #F4A460 (sandy brown)
|
|
- Song → #87CEEB (sky blue) or #DDA0DD (plum)
|
|
|
|
Fixed assignment:
|
|
- Malabar → #FFD700 (gold, already set)
|
|
|
|
Reply with ONLY a JSON object mapping speaker names to hex color codes:
|
|
{{"SpeakerName": "#RRGGBB", ...}}
|
|
|
|
JSON:"""
|
|
|
|
try:
|
|
response = client.chat.completions.create(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": "You assign colors to characters. Reply with ONLY valid JSON."},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
temperature=0.3,
|
|
max_tokens=500
|
|
)
|
|
|
|
message = response.choices[0].message
|
|
result = message.content or ""
|
|
|
|
# GLM models may put response in reasoning_content
|
|
if not result and hasattr(message, 'reasoning_content') and message.reasoning_content:
|
|
result = message.reasoning_content
|
|
|
|
# Try to parse JSON
|
|
json_match = re.search(r'\{[^}]+\}', result)
|
|
if json_match:
|
|
try:
|
|
parsed = json.loads(json_match.group())
|
|
for speaker, color in parsed.items():
|
|
if speaker in remaining_speakers:
|
|
color_mapping[speaker] = color
|
|
except json.JSONDecodeError:
|
|
print(f" Warning: Could not parse JSON response")
|
|
|
|
# Assign default hex colors for any remaining speakers
|
|
for speaker in remaining_speakers:
|
|
if speaker not in color_mapping:
|
|
# Simple fallback based on name (using hex codes)
|
|
name_lower = speaker.lower()
|
|
if 'mars' in name_lower:
|
|
color_mapping[speaker] = "#FF4500" # Orange red
|
|
elif 'earth' in name_lower:
|
|
color_mapping[speaker] = "#228B22" # Forest green
|
|
elif 'moon' in name_lower:
|
|
color_mapping[speaker] = "#C0C0C0" # Silver
|
|
elif 'sun' in name_lower:
|
|
color_mapping[speaker] = "#FFD700" # Gold
|
|
elif 'jupiter' in name_lower:
|
|
color_mapping[speaker] = "#D2691E" # Chocolate/orange
|
|
elif 'star' in name_lower:
|
|
color_mapping[speaker] = "#FFFFFF" # White
|
|
elif 'galaxy' in name_lower:
|
|
color_mapping[speaker] = "#9370DB" # Medium purple
|
|
elif 'volcano' in name_lower:
|
|
color_mapping[speaker] = "#8B0000" # Dark red
|
|
elif 'song' in name_lower:
|
|
color_mapping[speaker] = "#87CEEB" # Sky blue
|
|
else:
|
|
color_mapping[speaker] = "#808080" # Gray
|
|
|
|
return color_mapping
|
|
|
|
except Exception as e:
|
|
print(f" Error assigning colors: {e}")
|
|
# Return defaults for all remaining speakers
|
|
for speaker in remaining_speakers:
|
|
color_mapping[speaker] = "gray"
|
|
return color_mapping
|
|
|
|
|
|
def main():
|
|
# Get LLM config
|
|
base_url, model = get_llm_config()
|
|
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=base_url)
|
|
|
|
print(f"Using model: {model}")
|
|
print(f"Endpoint: {base_url or 'OpenAI default'}")
|
|
|
|
# Check input directory
|
|
if not INPUT_DIR.exists():
|
|
print(f"Error: Input directory {INPUT_DIR}/ not found")
|
|
sys.exit(1)
|
|
|
|
# Collect all speakers
|
|
print(f"\nCollecting speakers from {INPUT_DIR}/...")
|
|
speakers = collect_speakers(INPUT_DIR)
|
|
|
|
if not speakers:
|
|
print("Error: No speakers found")
|
|
sys.exit(1)
|
|
|
|
print(f"Found {len(speakers)} unique speakers:")
|
|
for speaker in sorted(speakers):
|
|
if speaker in FIXED_COLORS:
|
|
print(f" - {speaker}: {FIXED_COLORS[speaker]} (fixed)")
|
|
else:
|
|
print(f" - {speaker}")
|
|
|
|
# Assign colors
|
|
print(f"\nAssigning colors...")
|
|
color_mapping = assign_colors(speakers, client, model)
|
|
|
|
print(f"\nFinal color assignments:")
|
|
for speaker, color in sorted(color_mapping.items()):
|
|
fixed = " (fixed)" if speaker in FIXED_COLORS else ""
|
|
print(f" - {speaker}: {color}{fixed}")
|
|
|
|
# Save to JSON
|
|
with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
|
|
json.dump(color_mapping, f, ensure_ascii=False, indent=2)
|
|
|
|
print(f"\nSaved to: {OUTPUT_FILE}")
|
|
print(f"\nStep 5 Complete!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|