pensiev 0.25.5__py3-none-any.whl

memos/dataset_extractor_for_internvl2.py

@@ -0,0 +1,192 @@
+"""
+准备 image path 和 ocr 数据，用于后续提供给 internvl2 生成更符合需求的 caption
+"""
+
+import json
+import argparse
+from sqlalchemy.orm import sessionmaker
+from memos.models import EntityModel, EntityMetadataModel
+from memos.config import get_database_path
+from sqlalchemy import create_engine
+from tqdm import tqdm
+from pathlib import Path
+import argilla as rg
+from PIL import Image
+import io
+
+
+def prepare_huggingface_dataset(output_file, batch_size=100, record_count=10000, libraries=None):
+    """Prepare a Hugging Face dataset and save it as JSONL."""
+    db_path = get_database_path()
+    engine = create_engine(f"sqlite:///{db_path}")
+    Session = sessionmaker(bind=engine)
+
+    def simplify_ocr_data(ocr_data):
+        simplified_data = []
+        for item in ocr_data:
+            # 提取左上角和右下角的坐标
+            dt_boxes = [item["dt_boxes"][0], item["dt_boxes"][2]]
+            # 保留 rec_txt 字段，剔除 score
+            simplified_data.append({
+                "dt_boxes": dt_boxes,
+                "rec_txt": item["rec_txt"]
+            })
+        return simplified_data
+
+    with Session() as session, open(output_file, "w", encoding="utf-8") as f:
+        query = session.query(EntityModel)
+        
+        # 如果指定了 libraries，添加过滤条件
+        if libraries:
+            query = query.filter(EntityModel.library_id.in_(libraries))
+        
+        total = query.count()
+
+        progress_bar = tqdm(
+            total=min(total, record_count), desc="Processing entities", unit="entity"
+        )
+        inserted_records = 0
+
+        for offset in range(0, total, batch_size):
+            batch = query.limit(batch_size).offset(offset).all()
+
+            for entity in batch:
+                # Skip entities with "low_info" tag
+                if any(tag.name == "low_info" for tag in entity.tags):
+                    progress_bar.update(1)
+                    continue
+
+                metadata = {entry.key: entry.value for entry in entity.metadata_entries}
+
+                ocr = metadata.get("ocr_result")
+                if not ocr or not Path(entity.filepath).exists():
+                    progress_bar.update(1)
+                    continue
+
+                # Load JSON from string
+                ocr_data = json.loads(ocr)
+
+                # Simplify OCR result
+                simplified_ocr = simplify_ocr_data(ocr_data)
+
+                record = {
+                    "id": entity.id,
+                    "image": entity.filepath,
+                    "ocr": simplified_ocr,
+                }
+
+                # Dump to JSON string
+                json_record = json.dumps(record, ensure_ascii=False)
+                f.write(json_record + "\n")
+                progress_bar.update(1)
+                inserted_records += 1
+
+                if inserted_records >= record_count:
+                    break
+            if inserted_records >= record_count:
+                break
+
+        progress_bar.close()
+
+    print(f"Dataset saved to {output_file}")
+
+
+def init_argilla_dataset(client, dataset_name="image_captioning"):
+    workspace_name = "argilla"
+
+    workspace = client.workspaces(workspace_name)
+
+    if workspace is None:
+        workspace = rg.Workspace(name=workspace_name, client=client)
+        workspace.create()
+        print(f"Workspace created: {workspace_name}")
+
+    dataset = client.datasets(name=dataset_name)
+
+    if dataset is not None:
+        return dataset
+
+    settings = rg.Settings(
+        fields=[
+            rg.ImageField(name="image"), 
+            rg.TextField(name="filepath")
+        ],
+        questions=[
+            rg.TextQuestion(
+                name="text",
+                title="Description of the image",
+                required=True,
+                use_markdown=True,
+            )
+        ],
+    )
+
+    dataset = rg.Dataset(
+        name=dataset_name, workspace=workspace_name, settings=settings, client=client
+    )
+
+    dataset.create()
+    print(f"Dataset created: {dataset_name}")
+
+    return dataset
+
+
+def upload_to_argilla(input_file, batch_size=10, dataset_name="image_captioning"):
+    """Upload a JSONL dataset to Argilla."""
+
+    client = rg.Argilla(api_url="http://localhost:6900", api_key="argilla.apikey")
+
+    dataset = init_argilla_dataset(client, dataset_name)
+
+    records = []
+    total_records = sum(1 for _ in open(input_file, "r"))
+
+    with open(input_file, "r", encoding="utf-8") as f:
+        progress_bar = tqdm(
+            total=total_records, desc="Uploading to Argilla", unit="record"
+        )
+
+        for line in f:
+            record_data = json.loads(line)
+            image = Image.open(record_data["image"]).convert("RGB")
+            image.thumbnail((1280, 1280))
+
+            rg_record = rg.Record(
+                id=str(record_data["id"]),
+                fields={
+                    "image": image,
+                    "filepath": record_data["image"],
+                },
+                suggestions=[
+                    rg.Suggestion(
+                        "text", record_data["answer"], score=1.0, agent="internvl2"
+                    )
+                ],
+            )
+            records.append(rg_record)
+
+            if len(records) >= batch_size:
+                dataset.records.log(records)
+                progress_bar.update(batch_size)
+                records = []
+
+        if records:
+            dataset.records.log(records)
+            progress_bar.update(len(records))
+
+        progress_bar.close()
+
+    print(f"Dataset uploaded to Argilla: {dataset_name}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Prepare and upload dataset")
+    parser.add_argument("--output_file", default="dataset.jsonl", help="Output file path")
+    parser.add_argument("--size", type=int, default=10000, help="Number of records to extract")
+    parser.add_argument("--libraries", nargs="+", type=int, help="List of library IDs to filter entities")
+    args = parser.parse_args()
+
+    prepare_huggingface_dataset(args.output_file, record_count=args.size, libraries=args.libraries)
+    print(f"Dataset saved to {args.output_file}")
+    # Uncomment the following line if you want to upload to Argilla
+    # upload_to_argilla(args.output_file)

memos/default_config.yaml

@@ -0,0 +1,88 @@
+base_dir: ~/.memos
+
+# Database settings
+# Can be either:
+# 1. A file path for SQLite (relative to base_dir): database.db
+# 2. SQLite URL: sqlite:///absolute/path/to/database.db
+# 3. PostgreSQL URL: postgresql://postgres:mysecretpassword@localhost:5432/postgres
+database_path: database.db
+
+default_library: screenshots
+screenshots_dir: screenshots
+
+server_host: 0.0.0.0
+server_port: 8839
+
+# Enable authentication by uncommenting the following lines
+# auth_username: admin
+# auth_password: changeme
+
+default_plugins:
+- builtin_ocr
+# - builtin_vlm
+
+# using ollama as the vlm server
+vlm:
+  concurrency: 8
+  endpoint: http://localhost:11434
+  force_jpeg: true
+  modelname: minicpm-v
+  # 中文版本
+  prompt: 请帮描述这个图片中的内容，包括画面格局、出现的视觉元素等
+  # English version
+  prompt: Please describe the content of this image, including the layout and visual elements.
+  token: ''
+
+# using local ocr
+ocr:
+  concurrency: 8
+  # this is not used if use_local is true
+  endpoint: http://localhost:5555/predict
+  force_jpeg: false
+  token: ''
+  use_local: true
+
+# using local embedding for English as the main language
+embedding:
+  model: arkohut/jina-embeddings-v2-base-en
+  num_dim: 768
+  use_local: true
+  use_modelscope: false
+
+watch:
+  # number of recent events to consider when calculating processing rates
+  rate_window_size: 10
+  # sparsity factor for file processing
+  # a higher value means less frequent processing
+  # 1.0 means process every file
+  sparsity_factor: 3.0
+  # initial processing interval for file processing, means process one file with plugins for every N files
+  # but will be adjusted automatically based on the processing rate
+  # 1 means process the first file at the beginning
+  processing_interval: 12
+
+# A watch config like this means process every file with plugins at the beginning
+# but if the processing rate is slower than file generated, the processing interval 
+# will be increased automatically
+# watch:
+#   rate_window_size: 10
+#   sparsity_factor: 1.0
+#   processing_interval: 1
+
+# using local embedding for Chinese as the main language
+# embedding:
+#   model: arkohut/jina-embeddings-v2-base-zh
+#   num_dim: 768
+#   use_local: true
+#   use_modelscope: true
+
+# using ollama embedding
+# embedding:
+#   endpoint: http://localhost:11434/v1/embeddings # this is not used
+#   model: arkohut/gte-qwen2-1.5b-instruct:q8_0
+#   num_dim: 1536
+#   use_local: false
+#   use_modelscope: false
+
+record_interval: 4 # seconds
+facet: false # support facet filter

memos/embedding.py

@@ -0,0 +1,129 @@
+from typing import List
+import numpy as np
+from .config import settings
+import logging
+import httpx
+import logfire
+from functools import lru_cache
+import hashlib
+import json
+
+# Configure logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables
+model = None
+device = None
+
+
+def init_embedding_model():
+    import torch
+    from sentence_transformers import SentenceTransformer
+    
+    global model, device
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    if settings.embedding.use_modelscope:
+        from modelscope import snapshot_download
+        model_dir = snapshot_download(settings.embedding.model)
+        logger.info(f"Model downloaded from ModelScope to: {model_dir}")
+    else:
+        model_dir = settings.embedding.model
+        logger.info(f"Using model: {model_dir}")
+
+    model = SentenceTransformer(model_dir, trust_remote_code=True)
+    model.to(device)
+    logger.info(f"Embedding model initialized on device: {device}")
+
+
+def generate_embeddings(texts: List[str]) -> List[List[float]]:
+    global model
+
+    if model is None:
+        init_embedding_model()
+
+    if not texts:
+        return []
+
+    embeddings = model.encode(texts, convert_to_tensor=True, show_progress_bar=False)
+    embeddings = embeddings.cpu().numpy()
+
+    # Normalize embeddings
+    norms = np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)
+    norms[norms == 0] = 1
+    embeddings = embeddings / norms
+
+    return embeddings.tolist()
+
+
+def _hash_texts(texts: List[str]) -> str:
+    """Generate a stable hash for a list of texts."""
+    texts_json = json.dumps(texts, sort_keys=True)
+    return hashlib.sha256(texts_json.encode()).hexdigest()
+
+
+@logfire.instrument
+@lru_cache(maxsize=100)  # Cache last 100 requests
+def get_embeddings_cached(texts_hash: str, texts_tuple: tuple) -> List[List[float]]:
+    """Internal cached function that works with immutable types."""
+    texts = list(texts_tuple)
+    if settings.embedding.use_local:
+        embeddings = generate_embeddings(texts)
+    else:
+        embeddings = get_remote_embeddings(texts)
+
+    # Round the embedding values to 5 decimal places
+    return [
+        [round(float(x), 5) for x in embedding]
+        for embedding in embeddings
+    ]
+
+
+@logfire.instrument
+def get_embeddings(texts: List[str]) -> List[List[float]]:
+    """Get embeddings with caching support."""
+    # Convert texts to immutable type and create a hash
+    texts_hash = _hash_texts(texts)
+    texts_tuple = tuple(texts)
+    return get_embeddings_cached(texts_hash, texts_tuple)
+
+
+def get_remote_embeddings(texts: List[str]) -> List[List[float]]:
+    headers = {
+        "Content-Type": "application/json"
+    }
+    
+    if settings.embedding.token.get_secret_value():
+        headers["Authorization"] = f"Bearer {settings.embedding.token.get_secret_value()}"
+
+    endpoint = settings.embedding.endpoint
+    is_ollama = endpoint.endswith("/embed")
+
+    if is_ollama:
+        payload = {"model": settings.embedding.model, "input": texts}
+    else:  # openai compatible api
+        payload = {
+            "input": texts,
+            "model": settings.embedding.model,
+            "encoding_format": "float"
+        }
+
+    with httpx.Client(timeout=60) as client:
+        try:
+            response = client.post(endpoint, json=payload, headers=headers)
+            response.raise_for_status()
+            result = response.json()
+
+            if is_ollama:
+                return result["embeddings"]
+            else:  # openai compatible api
+                return [item["embedding"] for item in result["data"]]
+        except httpx.RequestError as e:
+            logger.error(f"Error fetching embeddings from remote endpoint: {e}")
+            return []  # Return an empty list instead of raising an exception

memos/frame_extractor.py

@@ -0,0 +1,53 @@
+import cv2
+from PIL import Image
+from pathlib import Path
+import argparse
+
+
+def extract_video_frame(video_path: Path, frame_number: int) -> Image.Image:
+    cap = cv2.VideoCapture(str(video_path))
+    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
+    ret, frame = cap.read()
+    cap.release()
+
+    if not ret:
+        return None
+
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(frame_rgb)
+
+
+def test_extract_video_frame(video_path: Path, frame_number: int):
+    # 确保测试视频文件存在
+    if not video_path.is_file():
+        print(f"Error: Video file not found at {video_path}")
+        return
+
+    # 尝试提取指定帧
+    frame_image = extract_video_frame(video_path, frame_number)
+
+    if frame_image is None:
+        print(f"Error: Failed to extract frame {frame_number} from video")
+    else:
+        print(f"Successfully extracted frame {frame_number} from video")
+        print(f"Frame dimensions: {frame_image.size}")
+
+        # 保存提取的帧为图像文件
+        output_path = video_path.with_name(f"extracted_frame_{frame_number}.png")
+        frame_image.save(output_path)
+        print(f"Extracted frame saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Extract a frame from a video file.")
+    parser.add_argument("video_path", type=str, help="Path to the video file")
+    parser.add_argument(
+        "frame_number", type=int, help="Frame number to extract (0-based index)"
+    )
+
+    args = parser.parse_args()
+
+    video_path = Path(args.video_path)
+    frame_number = args.frame_number
+
+    test_extract_video_frame(video_path, frame_number)

memos/logging_config.py

@@ -0,0 +1,35 @@
+import logging
+import sys
+
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "default": {
+            "format": "%(asctime)s - %(levelname)s - %(message)s",
+        },
+    },
+    "handlers": {
+        "default": {
+            "level": "INFO",
+            "formatter": "default",
+            "class": "logging.StreamHandler",
+            "stream": sys.stdout,
+        },
+    },
+    "loggers": {
+        "": {
+            "handlers": ["default"],
+            "level": "INFO",
+            "propagate": False,
+        },
+        "uvicorn.error": {
+            "level": "INFO",
+        },
+        "uvicorn.access": {
+            "handlers": ["default"],
+            "level": "INFO",
+            "propagate": False,
+        },
+    },
+}

memos/main.py

@@ -0,0 +1,104 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import List
+import time
+import random
+from fastapi import Response, HTTPException
+from datetime import datetime
+
+app = FastAPI()
+libraries = []
+
+
+class Folder(BaseModel):
+    path: str
+    libraryId: int
+
+
+class Library(BaseModel):
+    id: int
+    name: str
+    description: str | None
+    folders: List[Folder] = []
+    lastScanAt: datetime | None = None
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "id": 1,
+                    "name": "Main Library",
+                    "description": "A primary collection of various documents.",
+                    "folders": [
+                        {"path": "/documents/reports", "libraryId": 1},
+                        {"path": "/documents/notes", "libraryId": 1},
+                    ],
+                    "lastScanAt": "2023-10-04T14:48:00",
+                }
+            ]
+        }
+    }
+
+
+@app.get("/")
+async def root():
+    return "ok"
+
+
+@app.get("/libraries", response_model=List[Library])
+async def get_libraries():
+    return libraries
+
+
+@app.get("/libraries/{library_id}", response_model=Library)
+async def get_library(library_id: int):
+    for library in libraries:
+        if library.id == library_id:
+            return library
+    raise HTTPException(status_code=404, detail="Library not found")
+
+
+class LibraryParam(BaseModel):
+    name: str
+    description: str | None
+    folders: List[str]
+
+
+@app.post("/libraries", status_code=201)
+async def create_library(library: LibraryParam):
+    nextid = int(time.time()) + random.randint(1, 1000)
+    new_library = Library(
+        id=nextid,
+        name=library.name,
+        description=library.description,
+        folders=[Folder(path=path, libraryId=nextid) for path in library.folders],
+    )
+    libraries.append(new_library)
+    return new_library
+
+
+@app.put("/libraries/{library_id}")
+async def update_library(library_id: int, library: LibraryParam):
+    for lib in libraries:
+        if lib.id == library_id:
+            lib.name = library.name
+            lib.description = library.description
+            lib.folders = [
+                Folder(path=path, libraryId=library_id) for path in library.folders
+            ]
+            return Response(status_code=204)
+    raise HTTPException(status_code=404, detail="Library not found")
+
+
+@app.delete("/libraries/{library_id}", status_code=204)
+async def delete_library(library_id: int):
+    for lib in libraries:
+        if lib.id == library_id:
+            libraries.remove(lib)
+            return Response(status_code=204)
+    raise HTTPException(status_code=404, detail="Library not found")
+
+
+@app.post("/libraries/{library_id}/scan_tasks", status_code=202)
+async def request_scan_library(library_id):
+    pass

memos/migrations/alembic/README

@@ -0,0 +1 @@
+Generic single-database configuration.