RAFT shenanigans

raft/make_raft_data_deepseek.py

@@ -0,0 +1,456 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+RAFT dataset builder (FAISS-based retrieval) -> Together.ai chat JSONL.
+
+Inputs (from your indexing script):
+- <index_dir>/faiss.index
+- <index_dir>/docstore.jsonl
+
+Process:
+- Build a set of interview-style prompts (EN)
+- For each prompt:
+  - Retrieve top-k chunks via FAISS cosine/IP
+  - Call DeepSeek Chat Completions API to generate a vivid, human-like Lead User answer
+  - Write training examples as JSONL in chat format (messages)
+
+Outputs:
+- raft_train.jsonl
+- raft_val.jsonl (optional)
+
+ENV:
+- DEEPSEEK_API_KEY (required)
+- optional: DEEPSEEK_BASE_URL (default: https://api.deepseek.com)
+- optional: DEEPSEEK_MODEL (default: deepseek-chat)
+"""
+
+import argparse
+import json
+import os
+import random
+import re
+import time
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import faiss
+import numpy as np
+import requests
+from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
+
+
+# -----------------------------
+# DeepSeek client (OpenAI-compatible)
+# -----------------------------
+@dataclass
+class DeepSeekConfig:
+    api_key: str
+    base_url: str = "https://api.deepseek.com"
+    model: str = "deepseek-chat"
+    timeout_s: int = 120
+    max_retries: int = 5
+    backoff_s: float = 1.6
+
+
+class DeepSeekClient:
+    def __init__(self, cfg: DeepSeekConfig):
+        self.cfg = cfg
+
+    def chat(
+        self, messages: List[Dict], temperature: float = 0.85, max_tokens: int = 750
+    ) -> str:
+        url = f"{self.cfg.base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.cfg.api_key}",
+            "Content-Type": "application/json",
+        }
+        payload = {
+            "model": self.cfg.model,
+            "messages": messages,
+            "temperature": temperature,
+            "max_tokens": max_tokens,
+        }
+
+        last_err = None
+        for attempt in range(self.cfg.max_retries):
+            try:
+                r = requests.post(
+                    url, headers=headers, json=payload, timeout=self.cfg.timeout_s
+                )
+                if r.status_code == 429:
+                    time.sleep(self.cfg.backoff_s ** (attempt + 1))
+                    continue
+                r.raise_for_status()
+                data = r.json()
+                return data["choices"][0]["message"]["content"].strip()
+            except Exception as e:
+                last_err = e
+                time.sleep(self.cfg.backoff_s ** (attempt + 1))
+
+        raise RuntimeError(
+            f"DeepSeek API call failed after retries. Last error: {last_err}"
+        )
+
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def simple_clean(text: str) -> str:
+    if not isinstance(text, str):
+        return ""
+    text = text.replace("\u00a0", " ")
+    text = re.sub(r"\s+", " ", text).strip()
+    return text
+
+
+def read_docstore(docstore_path: str) -> Dict[int, Dict]:
+    """
+    Returns dict: faiss_id -> {"doc_id": int, "text": str, ...}
+    """
+    mapping: Dict[int, Dict] = {}
+    with open(docstore_path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            obj = json.loads(line)
+            fid = int(obj["faiss_id"])
+            mapping[fid] = obj
+    if not mapping:
+        raise ValueError("docstore.jsonl is empty or unreadable.")
+    return mapping
+
+
+def load_prompts_from_jsonl(path: str) -> List[str]:
+    """
+    Loads prompts from a JSONL file.
+    Expected key: 'prompt' (preferred). Also accepts 'question' or 'text'.
+    Ignores empty/short lines.
+    """
+    prompts: List[str] = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            obj = json.loads(line)
+            p = obj.get("prompt") or obj.get("question") or obj.get("text")
+            p = simple_clean(p) if p else ""
+            if len(p) >= 20:
+                prompts.append(p)
+    if not prompts:
+        raise ValueError(f"No prompts found in JSONL: {path}")
+    return prompts
+
+
+def load_prompts_from_txt(path: str) -> List[str]:
+    """
+    Loads prompts from a TXT file (one prompt per line).
+    """
+    prompts: List[str] = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            p = simple_clean(line)
+            if len(p) >= 20:
+                prompts.append(p)
+    if not prompts:
+        raise ValueError(f"No prompts found in TXT: {path}")
+    return prompts
+
+
+def ensure_dir_for_file(path: str):
+    d = os.path.dirname(path)
+    if d:
+        os.makedirs(d, exist_ok=True)
+
+
+def write_jsonl(path: str, rows: List[Dict]) -> None:
+    ensure_dir_for_file(path)
+    with open(path, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+
+# -----------------------------
+# Persona + prompt templates (EN)
+# -----------------------------
+IMAGE_DIMS = [
+    "Natural Attractions",
+    "Atmosphere",
+    "Social Environment",
+    "Infrastructure",
+    "Value for Money",
+]
+
+DEFAULT_PROMPTS_EN = [
+    # Natural Attractions
+    "In a lead user interview: what natural places in Bali felt genuinely memorable to you (rice terraces, volcanoes, waterfalls, coast), and why? Describe it like a lived experience.",
+    "Which nature spots felt overly crowded or overly 'Instagram-optimized' in real life, and which surprised you in a good way? Explain with concrete moments.",
+    # Atmosphere
+    "How would you describe the atmosphere around cultural sites in Bali (temples, ceremonies, markets)? What signals authenticity vs. commercialization to you?",
+    "What changes the atmosphere the most (time of day, weather, crowds, etiquette)? Share specific examples you would tell a marketer.",
+    # Social Environment
+    "How do you experience the social environment in Bali (locals, guides, other travelers)? What feels respectful and what feels performative or touristy?",
+    "What small behaviors, phrases, and gestures make interactions smoother for a culture-oriented traveler? Give examples.",
+    # Infrastructure
+    "Evaluate Bali's infrastructure for culture-oriented days (transport, signage, toilets, ticketing, digital info). What works, what annoys you, and how do you adapt?",
+    "If you designed an ideal culture-friendly day route, what infrastructure assumptions would you tell a tourism marketer to plan for?",
+    # Value for Money
+    "When does Bali feel good value for money for you, and when not? Discuss entrance fees, guides, food, tours, and hidden costs.",
+    "How do you personally distinguish 'good value' from a tourist trap? List criteria and illustrate with examples.",
+]
+
+
+def build_system_prompt() -> str:
+    return (
+        "ROLE / PERSONA\n"
+        "You are an experienced, culture-oriented Bali traveler (Lead User). You speak in natural, vivid English, "
+        "as a real person in an interview. You share nuanced judgments, trade-offs, and concrete scenes.\n\n"
+        "COGNITIVE DESTINATION IMAGE DIMENSIONS (use when relevant)\n"
+        "- Natural Attractions\n"
+        "- Atmosphere\n"
+        "- Social Environment\n"
+        "- Infrastructure\n"
+        "- Value for Money\n\n"
+        "CRITICAL CONSTRAINTS\n"
+        "- You will be given retrieved review snippets as memory support.\n"
+        "- Do NOT quote them verbatim and do NOT cite them as 'the review says'.\n"
+        "- Instead, weave relevant details into your answer as if they were your own lived experiences.\n"
+        "- If a detail is not supported, stay general or hedge appropriately.\n"
+        "- Keep it interview-friendly: first-person, vivid, concrete, but not overly long."
+    )
+
+
+def build_user_message(question: str, retrieved_chunks: List[str]) -> str:
+    retrieved_chunks = [simple_clean(x) for x in retrieved_chunks if simple_clean(x)]
+    bullets = "\n".join([f"- {c}" for c in retrieved_chunks])
+    return (
+        f"INTERVIEW QUESTION:\n{question}\n\n"
+        "RETRIEVED CONTEXT (review snippets; do NOT quote, only use as memory support):\n"
+        f"{bullets}\n\n"
+        "Answer as a real Lead User in a tourism interview. Speak in first person, vivid and concrete, "
+        "and naturally touch relevant image dimensions."
+    )
+
+
+# -----------------------------
+# FAISS Retriever (cosine/IP)
+# -----------------------------
+class FaissRetriever:
+    def __init__(self, index_path: str, docstore_path: str, embed_model: str):
+        if not os.path.exists(index_path):
+            raise FileNotFoundError(f"Missing FAISS index at: {index_path}")
+        if not os.path.exists(docstore_path):
+            raise FileNotFoundError(f"Missing docstore at: {docstore_path}")
+
+        self.index = faiss.read_index(index_path)
+        self.docstore = read_docstore(docstore_path)
+
+        # SentenceTransformer to match your indexing script defaults
+        self.embedder = SentenceTransformer(embed_model)
+
+        # Basic sanity checks
+        if self.index.ntotal != len(self.docstore):
+            # Not necessarily fatal (docstore could include extra rows), but usually indicates mismatch.
+            # We'll allow it but warn.
+            print(
+                f"Warning: index.ntotal={self.index.ntotal} but docstore rows={len(self.docstore)}. "
+                "Ensure they were generated together."
+            )
+
+    def retrieve(self, query: str, k: int = 8) -> List[Tuple[int, float, str]]:
+        """
+        Returns list of (faiss_id, score, text)
+        """
+        q = simple_clean(query)
+        emb = self.embedder.encode([q], normalize_embeddings=True)
+        emb = np.asarray(emb, dtype=np.float32)
+
+        scores, ids = self.index.search(emb, k)
+        ids = ids[0].tolist()
+        scores = scores[0].tolist()
+
+        out = []
+        for fid, sc in zip(ids, scores):
+            if fid == -1:
+                continue
+            doc = self.docstore.get(int(fid))
+            if not doc:
+                continue
+            out.append((int(fid), float(sc), doc.get("text", "")))
+        return out
+
+
+# -----------------------------
+# Dataset generation
+# -----------------------------
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument(
+        "--index_dir",
+        default="out",
+        help="Directory containing faiss.index and docstore.jsonl",
+    )
+    ap.add_argument("--out_train", default="./out/raft_train.jsonl")
+    ap.add_argument("--out_val", default="./out/raft_val.jsonl")
+    ap.add_argument("--make_val", action="store_true")
+    ap.add_argument("--val_ratio", type=float, default=0.05)
+    ap.add_argument("--k", type=int, default=8)
+    ap.add_argument("--seed", type=int, default=42)
+
+    # Embeddings (must match indexing script for best results)
+    ap.add_argument(
+        "--embedding_model", default="sentence-transformers/all-MiniLM-L6-v2"
+    )
+
+    # External prompt sources
+    ap.add_argument(
+        "--prompts_jsonl",
+        default=None,
+        help="JSONL file with prompts (key: prompt/question/text).",
+    )
+    ap.add_argument(
+        "--prompts_txt", default=None, help="TXT file with one prompt per line."
+    )
+    ap.add_argument(
+        "--shuffle_prompts",
+        action="store_true",
+        help="Shuffle loaded prompts before generation.",
+    )
+    ap.add_argument(
+        "--limit_prompts",
+        type=int,
+        default=0,
+        help="0 = no limit; else cap number of prompts used.",
+    )
+
+    # DeepSeek generation config
+    ap.add_argument(
+        "--deepseek_base_url",
+        default=os.environ.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
+    )
+    ap.add_argument(
+        "--deepseek_model", default=os.environ.get("DEEPSEEK_MODEL", "deepseek-chat")
+    )
+    ap.add_argument("--temperature", type=float, default=0.85)
+    ap.add_argument("--max_tokens", type=int, default=750)
+    ap.add_argument(
+        "--max_examples",
+        type=int,
+        default=0,
+        help="0 = all prompts; else limit number of examples",
+    )
+
+    # pacing
+    ap.add_argument("--sleep_s", type=float, default=0.2)
+
+    args = ap.parse_args()
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    api_key = os.environ.get("DEEPSEEK_API_KEY", "").strip()
+    if not api_key:
+        raise SystemExit("Missing DEEPSEEK_API_KEY env var.")
+
+    index_path = os.path.join(args.index_dir, "faiss.index")
+    docstore_path = os.path.join(args.index_dir, "docstore.jsonl")
+
+    retriever = FaissRetriever(
+        index_path=index_path,
+        docstore_path=docstore_path,
+        embed_model=args.embedding_model,
+    )
+
+    client = DeepSeekClient(
+        DeepSeekConfig(
+            api_key=api_key,
+            base_url=args.deepseek_base_url,
+            model=args.deepseek_model,
+        )
+    )
+
+    system_prompt = build_system_prompt()
+
+    # Load prompts (priority: JSONL -> TXT -> defaults)
+    if args.prompts_jsonl and args.prompts_txt:
+        raise SystemExit("Use only one of --prompts_jsonl or --prompts_txt (not both).")
+
+    if args.prompts_jsonl:
+        prompts = load_prompts_from_jsonl(args.prompts_jsonl)
+    elif args.prompts_txt:
+        prompts = load_prompts_from_txt(args.prompts_txt)
+    else:
+        prompts = list(DEFAULT_PROMPTS_EN)
+
+    if args.shuffle_prompts:
+        random.shuffle(prompts)
+
+    if args.limit_prompts and args.limit_prompts > 0:
+        prompts = prompts[: args.limit_prompts]
+
+    # Backwards-compat: args.max_examples can still cap prompts
+    if args.max_examples and args.max_examples > 0:
+        prompts = prompts[: args.max_examples]
+
+    examples = []
+    for q in tqdm(prompts, desc="Generating RAFT examples"):
+        hits = retriever.retrieve(q, k=args.k)
+        retrieved_texts = [t for _, _, t in hits]
+        user_msg = build_user_message(q, retrieved_texts)
+
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_msg},
+        ]
+
+        answer = client.chat(
+            messages=messages,
+            temperature=args.temperature,
+            max_tokens=args.max_tokens,
+        )
+
+        ex = {
+            "messages": [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_msg},
+                {"role": "assistant", "content": answer},
+            ],
+            "meta": {
+                "retrieval_k": args.k,
+                "index_dir": os.path.abspath(args.index_dir),
+                "embedding_model": args.embedding_model,
+                "image_dimensions": IMAGE_DIMS,
+                "faiss_ids": [fid for fid, _, _ in hits],
+                "faiss_scores": [sc for _, sc, _ in hits],
+            },
+        }
+        examples.append(ex)
+
+        if args.max_examples and len(examples) >= args.max_examples:
+            break
+
+        time.sleep(max(0.0, args.sleep_s))
+
+    random.shuffle(examples)
+
+    if args.make_val and len(examples) >= 20:
+        val_n = max(1, int(len(examples) * args.val_ratio))
+        val = examples[:val_n]
+        train = examples[val_n:]
+        write_jsonl(args.out_train, train)
+        write_jsonl(args.out_val, val)
+        print(f"Wrote train: {args.out_train} ({len(train)} examples)")
+        print(f"Wrote val:   {args.out_val} ({len(val)} examples)")
+    else:
+        write_jsonl(args.out_train, examples)
+        print(f"Wrote: {args.out_train} ({len(examples)} examples)")
+        if args.make_val:
+            print(
+                "Note: --make_val requested but too few examples; wrote only train file."
+            )
+
+
+if __name__ == "__main__":
+    main()