add: ner python script

1 files changed, 171 insertions, 0 deletions

diff --git a/2025/2025-09-07-ner.py b/2025/2025-09-07-ner.py
new file mode 100644
index 0000000..1af9aea
--- /dev/null
+++ b/2025/2025-09-07-ner.py
@@ -0,0 +1,171 @@
+import json
+import html2text
+import re
+import sqlite3
+from pathlib import Path
+
+from mlx_lm.utils import load  # type: ignore
+from mlx_lm.generate import generate  # type: ignore
+
+model, tokenizer = load("mlx-community/SmolLM3-3B-8bit")
+
+system_prompt = """/no_think
+You are a JSON generator.
+
+You are a JSON generator.
+
+Extract ALL people, places, and topics from the provided input.
+
+Please use structured JSON for the output:
+
+```json
+{
+  "people": [ "string", "string", etc. ],
+  "places": [ "string", "string", etc. ],
+  "topics": [ "string", "string", etc. ],
+}
+```
+
+Return only JSON inside these markers:
+<json>
+{ ... }
+</json>
+
+Important: Your entire output must be a single valid JSON object.
+Do not add comments, explanations, or text outside the JSON braces.
+If you cannot comply, output {}.
+"""
+
+def init_database(db_path="posts_processed.db"):
+    """Initialize SQLite database with posts table"""
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    # Create table if it doesn't exist
+    cursor.execute('''
+        CREATE TABLE IF NOT EXISTS processed_posts (
+            post_id TEXT PRIMARY KEY,
+            json_data TEXT NOT NULL,
+            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+        )
+    ''')
+
+    conn.commit()
+    return conn
+
+def post_exists(conn, post_id):
+    """Check if a post_id already exists in the database"""
+    cursor = conn.cursor()
+    cursor.execute('SELECT 1 FROM processed_posts WHERE post_id = ?', (post_id,))
+    return cursor.fetchone() is not None
+
+def save_post(conn, post_id, json_data):
+    """Save processed post to database"""
+    cursor = conn.cursor()
+    cursor.execute('''
+        INSERT INTO processed_posts (post_id, json_data)
+        VALUES (?, ?)
+    ''', (post_id, json.dumps(json_data)))
+    conn.commit()
+
+def extract_json_from_response(response):
+    """Extract JSON from the model response between <json> markers"""
+    try:
+        # Look for content between <json> and </json> markers
+        import re
+        json_match = re.search(r'<json>\s*(.*?)\s*</json>', response, re.DOTALL)
+        if json_match:
+            json_str = json_match.group(1)
+            return json.loads(json_str)
+        else:
+            # If no markers found, try to parse the entire response as JSON
+            # First, try to find a JSON object in the response
+            json_match = re.search(r'\{.*\}', response, re.DOTALL)
+            if json_match:
+                return json.loads(json_match.group(0))
+    except json.JSONDecodeError as e:
+        print(f"Failed to parse JSON: {e}")
+    except Exception as e:
+        print(f"Error extracting JSON: {e}")
+
+    return {}
+
+def load_all_posts(posts_dir="posts"):
+    """Load all posts from all JSON files into a single list"""
+    all_posts = []
+
+    posts_path = Path(posts_dir)
+
+    for json_file in posts_path.glob("*.json"):
+        try:
+            with open(json_file, 'r', encoding='utf-8') as f:
+                posts = json.load(f)
+                all_posts.extend(posts)
+        except Exception as e:
+            print(f"Error reading {json_file}: {e}")
+
+    return all_posts
+
+def main():
+    # Initialize database
+    conn = init_database()
+
+    # Load posts
+    posts = load_all_posts()
+    total_posts = len(posts)
+
+    processed_count = 0
+    skipped_count = 0
+
+    for i, post in enumerate(posts, 1):
+        post_id = post['post_id']
+
+        # Skip if post already exists in database
+        if post_exists(conn, post_id):
+            skipped_count += 1
+            print(f"\rProgress: {i}/{total_posts} | Processed: {processed_count} | Skipped: {skipped_count}", end="", flush=True)
+            continue
+
+        # Extract and clean text
+        html = post['post_html']
+        if html is None:
+            html = "<body></body>"
+
+        text2 = html2text.html2text(html)
+        cleaned = re.sub(r'https?://\S+|www\.\S+', '', text2)
+
+        prompt = f"""<input>
+        {cleaned}
+        </input>
+        """
+
+        messages = [
+          { "role": "system", "content": system_prompt },
+          { "role": "user", "content": prompt }
+        ]
+        prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+
+        response = generate(
+            model,
+            tokenizer,
+            prompt=prompt,
+            verbose=False,
+            max_tokens=200
+        )
+
+        # Extract JSON from response and save to database
+        extracted_json = extract_json_from_response(response)
+
+        if extracted_json:
+            save_post(conn, post_id, extracted_json)
+            processed_count += 1
+
+        print(f"\rProgress: {i}/{total_posts} | Processed: {processed_count} | Skipped: {skipped_count}", end="", flush=True)
+
+    # Close database connection
+    conn.close()
+
+    print(f"\n✓ Complete: Processed {processed_count} new posts, skipped {skipped_count} existing posts")
+
+if __name__ == "__main__":
+    main()