Duplicate from Kaomoji/coomfinity

Co-authored-by: Baker <Kaomoji@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,13 @@
+from python:3.10
+
+workdir /app
+
+copy requirements.txt .
+run pip install -r requirements.txt
+
+run mkdir /.cache && chmod 777 /.cache
+run mkdir .chroma && chmod 777 .chroma
+copy . .
+expose 7860
+
+cmd ["python", "server.py"]

README.md

@@ -0,0 +1,12 @@
+---
+title: Coomfinity
+emoji: 😭
+colorFrom: yellow
+colorTo: yellow
+sdk: docker
+pinned: false
+license: mit
+duplicated_from: Kaomoji/coomfinity
+---
+
+Put `https://kaomoji-coomfinity.hf.space/` as sillytavern-extras URL into sillytavern <- URL is different when you clone my space!

requirements.txt

@@ -0,0 +1,10 @@
+flask
+flask-cors
+colorama
+--extra-index-url https://download.pytorch.org/whl/cu117
+torch==2.0.0+cu117
+transformers==4.28.1
+sentence-transformers
+chromadb
+markdown
+Werkzeug

server.py

@@ -0,0 +1,157 @@
+from flask import (
+    Flask,
+    jsonify,
+    request,
+    render_template_string,
+    abort,
+)
+from flask_cors import CORS
+import unicodedata
+import markdown
+import time
+import os
+import gc
+import base64
+from io import BytesIO
+from random import randint
+import hashlib
+from colorama import Fore, Style, init as colorama_init
+import chromadb
+import posthog
+from chromadb.config import Settings
+from sentence_transformers import SentenceTransformer
+from werkzeug.middleware.proxy_fix import ProxyFix
+
+colorama_init()
+
+port = 7860
+host = "0.0.0.0"
+
+embedding_model = 'sentence-transformers/all-mpnet-base-v2'
+
+print("Initializing ChromaDB")
+
+# disable chromadb telemetry
+posthog.capture = lambda *args, **kwargs: None
+chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False))
+chromadb_embedder = SentenceTransformer(embedding_model)
+chromadb_embed_fn = chromadb_embedder.encode
+
+# Flask init
+app = Flask(__name__)
+CORS(app)  # allow cross-domain requests
+app.config["MAX_CONTENT_LENGTH"] = 100 * 1024 * 1024
+
+app.wsgi_app = ProxyFix(
+    app.wsgi_app, x_for=2, x_proto=1, x_host=1, x_prefix=1
+)
+
+def get_real_ip():
+    return request.remote_addr
+
+@app.route("/", methods=["GET"])
+def index():
+    with open("./README.md", "r", encoding="utf8") as f:
+        content = f.read()
+    return render_template_string(markdown.markdown(content, extensions=["tables"]))
+
+
+@app.route("/api/modules", methods=["GET"])
+def get_modules():
+    return jsonify({"modules": ['chromadb']})
+
+@app.route("/api/chromadb", methods=["POST"])
+def chromadb_add_messages():
+    data = request.get_json()
+    if "chat_id" not in data or not isinstance(data["chat_id"], str):
+        abort(400, '"chat_id" is required')
+    if "messages" not in data or not isinstance(data["messages"], list):
+        abort(400, '"messages" is required')
+
+    ip = get_real_ip()
+    chat_id_md5 = hashlib.md5(f'{ip}-{data["chat_id"]}'.encode()).hexdigest()
+    collection = chromadb_client.get_or_create_collection(
+        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+    )
+
+    documents = [m["content"] for m in data["messages"]]
+    ids = [m["id"] for m in data["messages"]]
+    metadatas = [
+        {"role": m["role"], "date": m["date"], "meta": m.get("meta", "")}
+        for m in data["messages"]
+    ]
+
+    if len(ids) > 0:
+        collection.upsert(
+            ids=ids,
+            documents=documents,
+            metadatas=metadatas,
+        )
+
+    return jsonify({"count": len(ids)})
+
+
+@app.route("/api/chromadb/query", methods=["POST"])
+def chromadb_query():
+    data = request.get_json()
+    if "chat_id" not in data or not isinstance(data["chat_id"], str):
+        abort(400, '"chat_id" is required')
+    if "query" not in data or not isinstance(data["query"], str):
+        abort(400, '"query" is required')
+
+    if "n_results" not in data or not isinstance(data["n_results"], int):
+        n_results = 1
+    else:
+        n_results = data["n_results"]
+
+    ip = get_real_ip()
+    chat_id_md5 = hashlib.md5(f'{ip}-{data["chat_id"]}'.encode()).hexdigest()
+    collection = chromadb_client.get_or_create_collection(
+        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+    )
+
+    n_results = min(collection.count(), n_results)
+
+    messages = []
+    if n_results > 0:
+        query_result = collection.query(
+            query_texts=[data["query"]],
+            n_results=n_results,
+        )
+    
+        documents = query_result["documents"][0]
+        ids = query_result["ids"][0]
+        metadatas = query_result["metadatas"][0]
+        distances = query_result["distances"][0]
+    
+        messages = [
+            {
+                "id": ids[i],
+                "date": metadatas[i]["date"],
+                "role": metadatas[i]["role"],
+                "meta": metadatas[i]["meta"],
+                "content": documents[i],
+                "distance": distances[i],
+            }
+            for i in range(len(ids))
+        ]
+
+    return jsonify(messages)
+    
+@app.route("/api/chromadb/purge", methods=["POST"])
+def chromadb_purge():
+    data = request.get_json()
+    if "chat_id" not in data or not isinstance(data["chat_id"], str):
+        abort(400, '"chat_id" is required')
+
+    ip = get_real_ip()
+    chat_id_md5 = hashlib.md5(f'{ip}-{data["chat_id"]}'.encode()).hexdigest()
+    collection = chromadb_client.get_or_create_collection(
+        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+    )
+
+    deleted = collection.delete()
+    print("ChromaDB embeddings deleted", len(deleted))
+
+    return 'Ok', 200
+app.run(host=host, port=port)