adapt paper schema and arxiv_ingest.py

Author: MoSchmidt

backend/app/models/paper.py

@@ -3,8 +3,17 @@
 from datetime import date, datetime
 from typing import TYPE_CHECKING, List, Optional
 
-from sqlalchemy import JSON, BigInteger, Date, DateTime, String, Text
+from pgvector.sqlalchemy import Vector
 from sqlalchemy import Enum as SqlEnum
+from sqlalchemy import (
+    JSON,
+    BigInteger,
+    Date,
+    DateTime,
+    String,
+    Text,
+    Index,
+)
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 from sqlalchemy.sql import func
 
@@ -13,38 +22,68 @@
 
 
 class Paper(Base):
-    """Database representation of a scholarly paper."""
+    """Database representation of a scholarly paper with vector embeddings."""
 
     __tablename__ = "paper"
 
+    # ------------------
+    # Columns
+    # ------------------
     paper_id: Mapped[int] = mapped_column(BigInteger, primary_key=True, index=True)
     doi: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)
+
     source: Mapped[PaperSource] = mapped_column(
         SqlEnum(PaperSource, name="paper_source"), nullable=False
     )
+
     paper_type: Mapped[PaperType] = mapped_column(
         SqlEnum(PaperType, name="paper_type"),
         nullable=False,
         default=PaperType.PREPRINT,
         server_default=PaperType.PREPRINT.value,
     )
+
     title: Mapped[str] = mapped_column(String(512), nullable=False)
     authors: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True)
     abstract: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
     published_at: Mapped[Optional[date]] = mapped_column(Date, nullable=True)
     pdf_url: Mapped[Optional[str]] = mapped_column(String(512), nullable=True)
     url: Mapped[Optional[str]] = mapped_column(String(512), nullable=True)
+
     fetched_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True), nullable=False, server_default=func.now()
     )
 
+    # ------------------
+    # Vector Embedding
+    # ------------------
+    embedding: Mapped[Optional[List[float]]] = mapped_column(
+        Vector(768), nullable=True
+    )
+
+    # ------------------
+    # Relationships
+    # ------------------
     project_links: Mapped[List["ProjectPaper"]] = relationship(
         "ProjectPaper", back_populates="paper", cascade="all, delete-orphan"
     )
+
     projects: Mapped[List["Project"]] = relationship(
         "Project", secondary="project_paper", back_populates="papers"
     )
 
+    # ------------------
+    # Indexes (HNSW)
+    # ------------------
+    __table_args__ = (
+        Index(
+            "ix_paper_embedding_hnsw",
+            "embedding",
+            postgresql_using="hnsw",
+            postgresql_with={"m": 16, "ef_search": 40, "lists": 100},
+        ),
+    )
+
 
 if TYPE_CHECKING:
     from .project import Project

backend/ingestion/arxiv_ingest.py

@@ -54,12 +54,13 @@ def process_batch(batch_buffer: List[Dict], embedder: Specter2Embedder) -> List[
 
 # pylint: disable=too-many-arguments, too-many-positional-arguments, too-many-locals
 def ingest_arxiv_metadata_to_parquet(
-    path: str,
-    out_dir: str,
-    embedder: Specter2Embedder,
-    limit: Optional[int] = None,
-    batch_size: int = 64,
-    shard_size: int = 5000,
+        path: str,
+        out_dir: str,
+        embedder: Specter2Embedder,
+        skip: int = 0,
+        limit: Optional[int] = None,
+        batch_size: int = 64,
+        shard_size: int = 5000,
 ) -> None:
     """Ingest a JSONL ArXiv dump and write Parquet shards with embeddings."""
     json_file = Path(path)
@@ -70,6 +71,8 @@ def ingest_arxiv_metadata_to_parquet(
         raise FileNotFoundError(f"Dataset not found: {json_file}")
 
     logger.info("Reading: %s", json_file)
+    if skip:
+        logger.info("Skipping first %s lines...", skip)
     logger.info("Writing shards: %s", out_path)
 
     batch_buffer: List[Dict] = []
@@ -78,7 +81,10 @@ def ingest_arxiv_metadata_to_parquet(
 
     with json_file.open("r", encoding="utf-8") as f:
         for i, line in enumerate(f):
-            if limit and i >= limit:
+            if i < skip:
+                continue
+
+            if limit and i >= skip + limit:
                 break
 
             try:
@@ -115,9 +121,9 @@ def ingest_arxiv_metadata_to_parquet(
     parser = argparse.ArgumentParser(description="Ingest ArXiv metadata into Parquet shards.")
     parser.add_argument("--path", type=str, required=True)
     parser.add_argument("--out", type=str, required=True)
+    parser.add_argument("--skip", type=int, default=0)
     parser.add_argument("--limit", type=int, default=None)
     parser.add_argument("--batch-size", type=int, default=64)
-    parser.add_argument("--shard-size", type=int, default=5000)
     parser.add_argument("--device", type=str, default=None)
 
     args = parser.parse_args()
@@ -128,7 +134,8 @@ def ingest_arxiv_metadata_to_parquet(
         path=args.path,
         out_dir=args.out,
         embedder=cli_embedder,
+        skip=args.skip,
         limit=args.limit,
         batch_size=args.batch_size,
-        shard_size=args.shard_size,
+        shard_size=5000
     )

backend/requirements.txt

@@ -6,6 +6,7 @@ uvicorn[standard]==0.38.0
 sqlalchemy==2.0.44
 psycopg2-binary==2.9.11
 alembic==1.17.1
+pgvector==0.4.1
 
 # --- Data validation & settings ---
 pydantic==2.12.3