generate_mlx_lm_datasets_no_sklearn.py

import requests
import json
import os
import random
import re

# Konfiguration
PAPERLESS_URL = "http://localhost:8010/api"
USERNAME = "paperless"  # Byt ut mot ditt Paperless-användarnamn
PASSWORD = "Privacy2025_"  # Byt ut mot ditt Paperless-lösenord
OUTPUT_DIR = "./data"
API_VERSION = "9"  # Senaste API-versionen per juni 2025
TRAIN_FILE = os.path.join(OUTPUT_DIR, "train.jsonl")
VALID_FILE = os.path.join(OUTPUT_DIR, "valid.jsonl")
TEST_FILE = os.path.join(OUTPUT_DIR, "test.jsonl")
TRAIN_RATIO = 0.8
VALID_RATIO = 0.1
TEST_RATIO = 0.1
MIN_EXAMPLES = 10  # Minsta antal exempel (dokument eller chunkar)
MIN_PER_SET = 2  # Minsta antal exempel i valid och test
MAX_WORDS_PER_CHUNK = 500  # Max ord per chunk för att undvika >2048 tokens

def get_api_token():
    """Hämta API-token från Paperless-ngx."""
    url = f"{PAPERLESS_URL}/token/"
    headers = {"Accept": f"application/json; version={API_VERSION}"}
    payload = {"username": USERNAME, "password": PASSWORD}
    
    try:
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        return response.json().get("token")
    except requests.exceptions.RequestException as e:
        print(f"Fel vid hämtning av token: {e}")
        return None

def fetch_documents(token, query=""):
    """Hämta dokument från Paperless-ngx och returnera text och metadata."""
    url = f"{PAPERLESS_URL}/documents/?full_perms=true"
    if query:
        url += f"&query={query}"
    headers = {
        "Authorization": f"Token {token}",
        "Accept": f"application/json; version={API_VERSION}"
    }
    documents = []
    
    while url:
        try:
            response = requests.get(url, headers=headers)
            response.raise_for_status()
            data = response.json()
            
            for doc in data["results"]:
                doc_data = {
                    "id": doc.get("id"),
                    "title": doc.get("title", ""),
                    "content": doc.get("content", ""),
                    "created": doc.get("created", ""),
                    "document_type": doc.get("document_type", None),
                    "tags": doc.get("tags", []),
                    "custom_fields": doc.get("custom_fields", []),
                }
                if doc_data["content"]:
                    documents.append(doc_data)
            
            url = data.get("next")
            print(f"Hämtade {len(data['results'])} dokument, nästa sida: {url}")
            
        except requests.exceptions.RequestException as e:
            print(f"Fel vid hämtning av dokument: {e}")
            break
    
    return documents

def split_content(content, max_words=MAX_WORDS_PER_CHUNK):
    """Dela upp innehåll i chunkar med max antal ord."""
    words = content.split()
    chunks = []
    current_chunk = []
    word_count = 0
    
    for word in words:
        current_chunk.append(word)
        word_count += 1
        if word_count >= max_words:
            chunks.append(" ".join(current_chunk))
            current_chunk = []
            word_count = 0
    
    if current_chunk:
        chunks.append(" ".join(current_chunk))
    
    return chunks

def format_metadata(doc, chunk_idx, total_chunks):
    """Formatera metadata för ett dokumentchunk."""
    metadata_parts = []
    if doc["title"]:
        metadata_parts.append(f"[Title: {doc['title']}]")
    if doc["created"]:
        metadata_parts.append(f"[Created: {doc['created']}]")
    if doc["document_type"]:
        metadata_parts.append(f"[Type: {doc['document_type']}]")
    if doc["tags"]:
        tags_str = ", ".join(str(tag) for tag in doc["tags"])
        metadata_parts.append(f"[Tags: {tags_str}]")
    if doc["custom_fields"]:
        fields = [f"{field['field']}={field['value']}" for field in doc["custom_fields"]]
        fields_str = ", ".join(fields)
        metadata_parts.append(f"[Custom Fields: {fields_str}]")
    if total_chunks > 1:
        metadata_parts.append(f"[Chunk: {chunk_idx+1}/{total_chunks}]")
    
    return " ".join(metadata_parts)

def create_examples(doc):
    """Skapa exempel genom att dela upp dokument i chunkar."""
    content = doc["content"].strip().replace("\n", " ")
    chunks = split_content(content, max_words=MAX_WORDS_PER_CHUNK)
    examples = []
    
    for idx, chunk in enumerate(chunks):
        metadata_str = format_metadata(doc, idx, len(chunks))
        text = f"{metadata_str} Text: {chunk}" if metadata_str else chunk
        examples.append({"text": text})
    
    return examples

def save_jsonl(examples, filename):
    """Spara exempel till en JSONL-fil i MLX-LM-format."""
    try:
        with open(filename, "w", encoding="utf-8") as f:
            for example in examples:
                f.write(json.dumps(example, ensure_ascii=False) + "\n")
        print(f"Sparade {len(examples)} exempel till {filename}")
    except IOError as e:
        print(f"Fel vid skrivning till fil: {e}")

def main():
    # Hämta API-token
    token = get_api_token()
    if not token:
        print("Kunde inte hämta API-token. Kontrollera användarnamn, lösenord och URL.")
        return
    
    # Hämta dokument
    documents = fetch_documents(token, query="")
    if not documents:
        print("Inga dokument hämtades.")
        return
    
    # Skapa exempel genom att dela upp dokument
    examples = []
    for doc in documents:
        examples.extend(create_examples(doc))
    
    if len(examples) < MIN_EXAMPLES:
        print(f"För få exempel: {len(examples)}. Behöver minst {MIN_EXAMPLES}.")
        return
    
    # Blanda och dela upp i träning, validering och test
    random.seed(42)
    random.shuffle(examples)
    total = len(examples)
    
    test_size = max(MIN_PER_SET, int(total * TEST_RATIO))
    valid_size = max(MIN_PER_SET, int(total * VALID_RATIO))
    train_size = total - valid_size - test_size
    
    if train_size < MIN_PER_SET:
        extra = MIN_PER_SET - train_size
        valid_size -= extra // 2
        test_size -= extra // 2
        train_size = MIN_PER_SET
    
    train_examples = examples[:train_size]
    valid_examples = examples[train_size:train_size + valid_size]
    test_examples = examples[train_size + valid_size:]
    
    print(f"Träning: {len(train_examples)} exempel, Validering: {len(valid_examples)} exempel, Test: {len(test_examples)} exempel")
    
    if len(valid_examples) < MIN_PER_SET or len(test_examples) < MIN_PER_SET:
        print(f"Fel: Validering ({len(valid_examples)}) eller Test ({len(test_examples)}) har för få exempel. Behöver minst {MIN_PER_SET} per set.")
        return
    
    # Spara till JSONL-filer
    save_jsonl(train_examples, TRAIN_FILE)
    save_jsonl(valid_examples, VALID_FILE)
    save_jsonl(test_examples, TEST_FILE)

if __name__ == "__main__":
    main()