raft-ft 2.2.0__py3-none-any.whl

raft/__init__.py

@@ -0,0 +1,9 @@
+"""RAFT: Retrieval-Augmented Fine-Tuning."""
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+__all__ = ["__version__"]
+
+__version__ = "2.2.0"

raft/chunker.py

@@ -0,0 +1,138 @@
+"""
+This module contains the chunker function, which splits blog posts into chunks
+of a maximum length.
+"""
+
+import json
+from typing import List, Dict, Tuple, Generator
+import tiktoken
+import pandas as pd
+
+MAX_EMBEDDING_LENGTH = 2048
+encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
+
+
+def map_line(line: str) -> Tuple[str, str, int]:
+    """
+    Encode and potentially truncate a line of text.
+
+    Args:
+        line (str): The input line of text.
+
+    Returns:
+        Tuple[str, str, int]: A tuple containing:
+            - the original or truncated line,
+            - the decoded version, and
+            - the length of the encoded line.
+    """
+    encoded = encoding.encode(line)
+    if len(encoded) < MAX_EMBEDDING_LENGTH:
+        return (line, encoding.decode(encoded), len(encoded))
+    else:
+        truncated_encoded = encoded[:MAX_EMBEDDING_LENGTH]
+        return (
+            line[:MAX_EMBEDDING_LENGTH],
+            encoding.decode(truncated_encoded),
+            MAX_EMBEDDING_LENGTH,
+        )
+
+
+def split_into_chunks(
+    blog_posts: pd.DataFrame,
+) -> Generator[Tuple[Dict[str, str], str], None, None]:
+    """
+    Split blog posts into chunks of a maximum length.
+
+    Args:
+        blog_posts (pd.DataFrame): DataFrame containing blog post data.
+
+    Yields:
+        Tuple[Dict[str, str], str]: A tuple containing metadata
+            and the chunked content.
+    """
+    for _, post in blog_posts.iterrows():
+        print(f"Splitting {post['title']}")
+        lines = list(map(map_line, post["content"].split("\n")))
+
+        # Calculate total parts
+        total_parts = 0
+        temp_chunk_length = 0
+        for _, _, length in lines:
+            if temp_chunk_length + length > MAX_EMBEDDING_LENGTH:
+                total_parts += 1
+                temp_chunk_length = length
+            else:
+                temp_chunk_length += length
+        total_parts += 1  # for the last chunk
+
+        chunk: List[str] = []
+        chunk_length = 0
+        part = 1
+
+        for line, _, length in lines:
+            if chunk_length + length > MAX_EMBEDDING_LENGTH:
+                metadata = {
+                    "title": post["title"],
+                    "url": post["link"],
+                    "date": (
+                        post["date"].isoformat()
+                        if isinstance(post["date"], pd.Timestamp)
+                        else post["date"]
+                    ),
+                    "total_parts": str(total_parts),
+                    "part": str(part),
+                }
+                yield metadata, "\n".join(chunk)
+                chunk = []
+                chunk_length = 0
+                part += 1
+            chunk.append(line)
+            chunk_length += length
+
+        if chunk:
+            metadata = {
+                "title": post["title"],
+                "url": post["link"],
+                "date": (
+                    post["date"].isoformat()
+                    if isinstance(post["date"], pd.Timestamp)
+                    else post["date"]
+                ),
+                "total_parts": str(total_parts),
+                "part": str(part),
+            }
+            yield metadata, "\n".join(chunk)
+
+
+def chunker(name: str) -> None:
+    """
+    Process a JSONL file of blog posts, split them into chunks,
+    and save the results.
+
+    Args:
+        name (str): The name of the file to process (without extension).
+    """
+    sourcefile = f"data/{name}.jsonl"
+    outputfile = f"data/{name}_chunked.jsonl"
+
+    try:
+        # Read the JSONL file line by line
+        with open(sourcefile, "r") as f:
+            data = [json.loads(line) for line in f]
+
+        blog_posts: pd.DataFrame = pd.DataFrame(data)
+        print(f"Splitting {len(blog_posts)} blog posts into chunks")
+
+        with open(outputfile, "w") as f:
+            for item in split_into_chunks(blog_posts):
+                f.write(json.dumps(item) + "\n")
+
+    except json.JSONDecodeError as e:
+        print(f"Error decoding JSON: {e}")
+        print(f"Error occurred at line {e.lineno}")
+        with open(sourcefile, "r") as f:
+            problematic_line = f.readlines()[e.lineno - 1]
+        print(f"Problematic line: {problematic_line}")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")

raft/cli.py

@@ -0,0 +1,189 @@
+"""
+This module contains the main CLI functionality for the RAFT project.
+"""
+
+import argparse
+from raft import (
+    files_helper,
+    embeddings_helpers,
+    substack_embeddings,
+    generate_finetune,
+    oai_finetune,
+    memories,
+)
+
+
+def action_doc() -> str:
+    """
+    Return a string describing available actions for the CLI.
+
+    Returns:
+        str: Description of available actions.
+    """
+    return """
+The following actions are available:
+
+- interactive: Guided end-to-end session: sources, conversations, finetune.
+- tweets: Build a dataset from tweets via ariadne interactive.
+- fetch: Fetch the blog from Substack and store it in the data directory.
+- chunk: Chunk the blog into 4096 token pieces and store them in /data.
+- embed: Create embeddings for the chunks and store them.
+- ft:gen: Generate finetune files for the blog.
+- ft:run: Run the finetune job (OpenAI, or huggingface via opbdh).
+- bench:setup: Setup the benchmark for the blog.
+- ask: Ask a question about the blog content.
+
+ft:run routes by --model: OpenAI-finetunable ids go to the OpenAI API,
+anything else (an org/name huggingface id) is trained on a GPU pod via
+opbdh; unknown flags after the name are forwarded verbatim to
+`opbdh launch` (e.g. --vram-gb 48 --max-spend 5).
+"""
+
+
+cmds = [
+    "interactive",
+    "tweets",
+    "fetch",
+    "chunk",
+    "embed",
+    "ft:gen",
+    "ft:run",
+    "bench:setup",
+    "ask",
+]
+
+
+def main() -> None:
+    """
+    Main function to handle CLI arguments and execute corresponding actions.
+    """
+    parser = argparse.ArgumentParser(
+        description="Run the raft command.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=action_doc(),
+    )
+    parser.add_argument(
+        "action",
+        help="The action to perform; see below for details.",
+        choices=cmds,
+    )
+
+    parser.add_argument(
+        "name",
+        nargs="?",
+        default="",
+        help="The name of the dataset to process.",
+    )
+    parser.add_argument(
+        "--oai",
+        help="Only generate finetune or benchmark for openai \
+            (from existing generic file).",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--generic",
+        help="Only generate generic finetune or benchmark file.",
+        default=False,
+        action="store_true",
+    )
+    parser.add_argument(
+        "--question",
+        help="The question to ask about the blog content.",
+        default=None,
+    )
+    parser.add_argument(
+        "--no-useful-check",
+        action="store_true",
+        help="Skip usefulness check when summarizing memories",
+    )
+    parser.add_argument(
+        "--model",
+        help="Model to finetune (ft:run): an OpenAI id, or a \
+            huggingface org/name id trained via opbdh.",
+        default="",
+    )
+    parser.add_argument(
+        "--no-interactive",
+        action="store_true",
+        help="ft:run: never prompt; rely on flags and opbdh config.",
+    )
+
+    args, extra = parser.parse_known_args()
+
+    if args.action not in ("ft:run",) and extra:
+        parser.error(f"unrecognized arguments: {' '.join(extra)}")
+
+    needs_name = args.action not in ("interactive", "tweets")
+    if needs_name and not args.name:
+        parser.error(f"the '{args.action}' action requires a dataset name")
+
+    if args.action == "interactive":
+        from .flows import run_interactive
+
+        run_interactive()
+    elif args.action == "tweets":
+        from .tweet_mode import run_tweet_mode
+
+        run_tweet_mode(args.name)
+    elif args.action == "fetch":
+        substack_embeddings.main(args.name)
+    elif args.action == "chunk":
+        files_helper.chunker(args.name)
+    elif args.action == "embed":
+        embeddings_helpers.store_grounding_embeddings(args.name)
+    elif args.action == "ft:gen":
+        if args.oai:
+            oai_finetune.create_openai_finetune_file(args.name)
+        elif args.generic:
+            generate_finetune.generate_finetune(args.name)
+        else:
+            generate_finetune.generate_finetune(args.name)
+            oai_finetune.create_openai_finetune_file(args.name)
+    elif args.action == "ft:run":
+        from .hf_finetune import is_openai_finetunable, run_hf_finetune
+        from .interactive import ask
+
+        model = args.model
+        if not model and not args.no_interactive:
+            model = ask(
+                "Model to finetune (OpenAI id or huggingface org/name)",
+                "gpt-4o-mini-2024-07-18",
+            )
+        if is_openai_finetunable(model):
+            if extra:
+                parser.error(
+                    "opbdh flags only apply to huggingface models: "
+                    f"{' '.join(extra)}"
+                )
+            oai_finetune.run_oai_finetune(args.name, model=model)
+        else:
+            run_hf_finetune(
+                args.name,
+                model,
+                opbdh_args=extra,
+                interactive=not args.no_interactive,
+            )
+    elif args.action == "bench:setup":
+        if args.oai:
+            oai_finetune.create_openai_finetune_file(args.name, "benchmark")
+        elif args.generic:
+            generate_finetune.generate_finetune(args.name)
+        else:
+            generate_finetune.generate_benchmark(args.name)
+            oai_finetune.create_openai_finetune_file(args.name, "benchmark")
+    elif args.action == "ask":
+        if args.question is None:
+            print("Please provide a question using the --question argument.")
+        else:
+            memory_manager = memories.MemoryManager(
+                args.name, {}
+            )  # Empty metadata for now
+            answer = memory_manager.ask_question(args.question)
+            print(f"Answer: {answer}")
+    else:
+        print(f"Unknown action: {args.action}")
+
+
+if __name__ == "__main__":
+    main()

raft/convo_structurer.py

@@ -0,0 +1,338 @@
+"""
+Turn conversation examples and text sources -- structured or not -- into
+the datasets raft consumes:
+
+- transcripts: data/{name}_transcript_{i}.json
+  {"participants": {"q": ..., "a": ...}, "date": ..., "url": ...,
+   "exchanges": [[question, answer], ...]}
+- grounding corpus: data/{name}.jsonl, one {"title", "link", "date",
+  "content"} object per line (the input of `raft chunk`).
+
+Unstructured files are converted with an LLM; structured ones (raft
+transcripts, chat-message arrays, grounding jsonl) are recognised and
+passed through.
+"""
+
+import json
+import os
+import re
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional, Tuple
+
+from openai import OpenAI
+
+from . import hx
+
+STRUCTURER_MODEL = os.environ.get("RAFT_LLM_MODEL", "gpt-4o")
+
+# Keep single LLM calls comfortably inside the context window.
+MAX_STRUCTURE_CHARS = 24000
+
+STRUCTURER_SYSTEM_PROMPT = """\
+You extract question/answer exchanges from a raw conversation record
+(interview, chat log, podcast transcript, email thread...).
+
+The target is the person being emulated: their utterances are answers,
+everyone else's are questions. Merge consecutive utterances by the same
+side. Drop utterances that are not part of a question/answer flow
+(stage directions, ads, timestamps).
+
+Reply with JSON only, in this exact shape:
+{
+  "participants": {"q": "<questioner name>", "a": "<target name>"},
+  "date": "<ISO date if inferable, else null>",
+  "url": "<source url if present, else null>",
+  "exchanges": [["<question>", "<answer>"], ...]
+}"""
+
+
+def transcript_path(name: str, index: int) -> str:
+    """Return the path of transcript #index for a dataset name."""
+    return f"data/{name}_transcript_{index}.json"
+
+
+def next_transcript_index(name: str) -> int:
+    """Return the first unused transcript index for a dataset name."""
+    i = 1
+    while os.path.exists(transcript_path(name, i)):
+        i += 1
+    return i
+
+
+def write_transcript(
+    name: str,
+    participants: Dict[str, str],
+    date: str,
+    url: str,
+    exchanges: List[List[str]],
+    index: Optional[int] = None,
+) -> str:
+    """
+    Write a transcript file in the format generate_finetune expects.
+
+    Returns:
+        str: The path written.
+    """
+    if index is None:
+        index = next_transcript_index(name)
+    path = transcript_path(name, index)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w") as f:
+        json.dump(
+            {
+                "participants": participants,
+                "date": date,
+                "url": url,
+                "exchanges": exchanges,
+            },
+            f,
+            indent=2,
+        )
+    return path
+
+
+def is_transcript(data: Any) -> bool:
+    """Check whether parsed JSON already matches the transcript schema."""
+    return (
+        isinstance(data, dict)
+        and isinstance(data.get("exchanges"), list)
+        and isinstance(data.get("participants"), dict)
+        and {"q", "a"} <= set(data["participants"])
+    )
+
+
+def parse_json_or_jsonl(raw: str) -> Any:
+    """
+    Parse a string as JSON, falling back to JSONL (returning the list of
+    parsed lines). Returns None if it is neither.
+    """
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        pass
+    try:
+        parsed = [json.loads(line) for line in raw.splitlines() if line.strip()]
+    except json.JSONDecodeError:
+        return None
+    return parsed or None
+
+
+def is_message_array(value: Any) -> bool:
+    """Check whether a value looks like a chat-message array."""
+    return (
+        isinstance(value, list)
+        and bool(value)
+        and all(isinstance(m, dict) for m in value)
+        and any("role" in m for m in value)
+    )
+
+
+def find_message_arrays(data: Any) -> List[List[Dict[str, Any]]]:
+    """
+    Collect every chat-message array in a parsed conversation file.
+
+    Handles a bare message array, {"messages": [...]}, ariadne's
+    {"conversations": [{"messages": [...]}, ...]}, and JSONL files whose
+    lines are any of the above.
+
+    Returns:
+        List[List[Dict[str, Any]]]: One entry per conversation found.
+    """
+    if is_message_array(data):
+        return [data]
+    if isinstance(data, dict):
+        if is_message_array(data.get("messages")):
+            return [data["messages"]]
+        if isinstance(data.get("conversations"), list):
+            found = []
+            for convo in data["conversations"]:
+                found.extend(find_message_arrays(convo))
+            return found
+        return []
+    if isinstance(data, list):
+        found = []
+        for item in data:
+            found.extend(find_message_arrays(item))
+        return found
+    return []
+
+
+def messages_to_exchanges(
+    messages: List[Dict[str, Any]], target: str
+) -> Tuple[List[List[str]], str]:
+    """
+    Convert a chat-message array (role/content, optionally name) into
+    [question, answer] exchanges, treating the target's (or assistant's)
+    turns as answers.
+
+    Returns:
+        Tuple[List[List[str]], str]: The exchanges and the most frequent
+        questioner name found (empty string if none).
+    """
+
+    def is_answer(msg: Dict[str, Any]) -> bool:
+        if target and msg.get("name", "").lower() == target.lower():
+            return True
+        return msg.get("role") == "assistant"
+
+    exchanges: List[List[str]] = []
+    q_names: Dict[str, int] = {}
+    pending_q: List[str] = []
+    for msg in messages:
+        content = str(msg.get("content", "")).strip()
+        if not content or msg.get("role") == "system":
+            continue
+        if is_answer(msg):
+            if pending_q:
+                exchanges.append(["\n".join(pending_q), content])
+                pending_q = []
+            elif exchanges:
+                exchanges[-1][1] += f"\n{content}"
+        else:
+            pending_q.append(content)
+            q_name = msg.get("name", "")
+            if q_name:
+                q_names[q_name] = q_names.get(q_name, 0) + 1
+
+    questioner = max(q_names, key=lambda k: q_names[k]) if q_names else ""
+    return exchanges, questioner
+
+
+def structure_raw_conversation(raw_text: str, target: str) -> Dict[str, Any]:
+    """
+    Use an LLM to turn an unstructured conversation into transcript JSON.
+
+    Args:
+        raw_text (str): The raw conversation text.
+        target (str): Name of the person being emulated.
+
+    Returns:
+        Dict[str, Any]: participants/date/url/exchanges dict.
+    """
+    client = OpenAI()
+    if len(raw_text) > MAX_STRUCTURE_CHARS:
+        hx.warn(
+            f"source is {len(raw_text)} chars; structuring the first "
+            f"{MAX_STRUCTURE_CHARS} -- split long files for full coverage"
+        )
+        raw_text = raw_text[:MAX_STRUCTURE_CHARS]
+    response = client.chat.completions.create(
+        model=STRUCTURER_MODEL,
+        response_format={"type": "json_object"},
+        messages=[
+            {"role": "system", "content": STRUCTURER_SYSTEM_PROMPT},
+            {
+                "role": "user",
+                "content": f"The target is: {target}\n\n---\n\n{raw_text}",
+            },
+        ],
+    )
+    data = json.loads(response.choices[0].message.content or "{}")
+    if not isinstance(data.get("exchanges"), list) or not data["exchanges"]:
+        raise ValueError("the LLM could not extract any exchanges")
+    data.setdefault("participants", {})
+    data["participants"].setdefault("q", "Interviewer")
+    data["participants"].setdefault("a", target)
+    return data
+
+
+def import_conversation_file(name: str, path: str, target: str) -> str:
+    """
+    Import one conversation file (structured or not) as a transcript.
+
+    Recognises: raft transcript JSON, chat-message arrays (bare list or
+    {"messages": [...]} / {"conversations": [...]}), anything else is
+    treated as raw text and structured with an LLM.
+
+    Returns:
+        str: The transcript path written.
+    """
+    with open(path) as f:
+        raw = f.read()
+
+    data = parse_json_or_jsonl(raw) if path.endswith((".json", ".jsonl")) else None
+    today = datetime.now(timezone.utc).date().isoformat()
+
+    if is_transcript(data):
+        return write_transcript(
+            name,
+            data["participants"],
+            data.get("date") or today,
+            data.get("url") or path,
+            data["exchanges"],
+        )
+
+    conversations = find_message_arrays(data)
+    if conversations:
+        exchanges: List[List[str]] = []
+        questioners: Dict[str, int] = {}
+        for messages in conversations:
+            got, questioner = messages_to_exchanges(messages, target)
+            exchanges.extend(got)
+            if questioner:
+                questioners[questioner] = questioners.get(questioner, 0) + 1
+        if not exchanges:
+            raise ValueError(f"{path}: no question/answer exchanges found")
+        top = max(questioners, key=lambda k: questioners[k]) if questioners else ""
+        meta = data if isinstance(data, dict) else {}
+        return write_transcript(
+            name,
+            {"q": top or "Interviewer", "a": target},
+            meta.get("date") or today,
+            meta.get("url") or path,
+            exchanges,
+        )
+
+    hx.step(f"{path}: unstructured, extracting exchanges with {STRUCTURER_MODEL}")
+    structured = structure_raw_conversation(raw, target)
+    return write_transcript(
+        name,
+        structured["participants"],
+        structured.get("date") or today,
+        structured.get("url") or path,
+        structured["exchanges"],
+    )
+
+
+def import_text_source_file(name: str, path: str) -> int:
+    """
+    Append a text source (structured grounding jsonl or raw text) to the
+    grounding corpus data/{name}.jsonl.
+
+    Returns:
+        int: Number of documents appended.
+    """
+    corpus_path = f"data/{name}.jsonl"
+    os.makedirs("data", exist_ok=True)
+
+    with open(path) as f:
+        raw = f.read()
+
+    records: List[Dict[str, Any]] = []
+    if path.endswith(".jsonl"):
+        try:
+            parsed = [json.loads(line) for line in raw.splitlines() if line.strip()]
+            if all(isinstance(p, dict) and "content" in p for p in parsed):
+                records = parsed
+        except json.JSONDecodeError:
+            records = []
+
+    if not records:
+        mtime = datetime.fromtimestamp(os.path.getmtime(path), timezone.utc)
+        title = re.sub(r"[_-]+", " ", os.path.splitext(os.path.basename(path))[0])
+        records = [
+            {
+                "title": title,
+                "link": path,
+                "date": mtime.date().isoformat(),
+                "content": raw,
+            }
+        ]
+
+    with open(corpus_path, "a") as f:
+        for record in records:
+            record.setdefault("title", "untitled")
+            record.setdefault("link", path)
+            record.setdefault("date", datetime.now(timezone.utc).date().isoformat())
+            f.write(json.dumps(record) + "\n")
+    return len(records)