Upload folder using huggingface_hub

b00c961 verified 11 months ago

3.3 kB

import os
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict
from sentence_transformers import SentenceTransformer
import uvicorn

app = FastAPI()

# --- State ---
SHARD_ID = os.getenv("SHARD_ID", "shard_0")
TOPIC = os.getenv("SHARD_TOPIC", "General")
print(f"🐝 [Shard {SHARD_ID}] Initializing... Topic: {TOPIC}")

# Load Embedding Model (Shared or Local)
# In production, shards might just do vector math if embeddings are pre-computed,
# but for standalone robustness, we load the model.
embedder = SentenceTransformer('all-MiniLM-L6-v2')

# Mock Knowledge Base (In prod, load from Disk/DB)
KNOWLEDGE_BASE = []

def seed_knowledge():
    print(f"🌱 Seeding Knowledge for {TOPIC}...")
    topics_data = {
        "Science": [
            "The speed of light is 299,792,458 m/s.",
            "Mitochondria is the powerhouse of the cell.",
            "Water boils at 100 degrees Celsius at sea level."
        ],
        "History": [
            "The Roman Empire fell in 476 AD.",
            "World War II ended in 1945.",
            "The Great Wall of China was built over centuries."
        ],
        "Coding": [
            "Python uses indentation for blocks.",
            "React uses a virtual DOM for performance.",
            "Docker containers share the host OS kernel."
        ]
    }
    
    facts = topics_data.get(TOPIC, ["Generic fact 1", "Generic fact 2"])
    
    for i, text in enumerate(facts):
        vec = embedder.encode(text)
        KNOWLEDGE_BASE.append({
            "id": f"{SHARD_ID}_{i}",
            "text": text,
            "vector": vec,
            "metadata": {"centrality": 0.9, "recency": 1.0, "weight": 1.0}
        })
    print(f"✅ Loaded {len(KNOWLEDGE_BASE)} facts.")

@app.on_event("startup")
async def startup_event():
    seed_knowledge()

# --- API ---

class RetrievalRequest(BaseModel):
    query_text: str
    query_vector: List[float] = None # Optional, if router computed it

class RetrievalResponse(BaseModel):
    shard_id: str
    best_text: str
    score: float

@app.post("/retrieve", response_model=RetrievalResponse)
async def retrieve(req: RetrievalRequest):
    # 1. Get Query Vector
    if req.query_vector:
        q_vec = np.array(req.query_vector)
    else:
        q_vec = embedder.encode(req.query_text)
        
    # 2. Vector Search (Dot Product)
    # Stack all vectors
    db_vecs = np.array([item["vector"] for item in KNOWLEDGE_BASE])
    
    # Cosine Sim
    norm_q = np.linalg.norm(q_vec)
    norm_db = np.linalg.norm(db_vecs, axis=1)
    dot = np.dot(db_vecs, q_vec)
    sims = dot / (norm_q * norm_db + 1e-9)
    
    # 3. Apply Metadata (P = S * C * R * W)
    best_score = -1.0
    best_idx = 0
    
    for i, sim in enumerate(sims):
        meta = KNOWLEDGE_BASE[i]["metadata"]
        p_score = sim * meta["centrality"] * meta["recency"] * meta["weight"]
        
        if p_score > best_score:
            best_score = p_score
            best_idx = i
            
    return RetrievalResponse(
        shard_id=SHARD_ID,
        best_text=KNOWLEDGE_BASE[best_idx]["text"],
        score=float(best_score)
    )

if __name__ == "__main__":
    port = int(os.getenv("PORT", 8001))
    uvicorn.run(app, host="0.0.0.0", port=port)

	import os
	import numpy as np
	from fastapi import FastAPI, HTTPException
	from pydantic import BaseModel
	from typing import List, Dict
	from sentence_transformers import SentenceTransformer
	import uvicorn

	app = FastAPI()

	# --- State ---
	SHARD_ID = os.getenv("SHARD_ID", "shard_0")
	TOPIC = os.getenv("SHARD_TOPIC", "General")
	print(f"🐝 [Shard {SHARD_ID}] Initializing... Topic: {TOPIC}")

	# Load Embedding Model (Shared or Local)
	# In production, shards might just do vector math if embeddings are pre-computed,
	# but for standalone robustness, we load the model.
	embedder = SentenceTransformer('all-MiniLM-L6-v2')

	# Mock Knowledge Base (In prod, load from Disk/DB)
	KNOWLEDGE_BASE = []

	def seed_knowledge():
	print(f"🌱 Seeding Knowledge for {TOPIC}...")
	topics_data = {
	"Science": [
	"The speed of light is 299,792,458 m/s.",
	"Mitochondria is the powerhouse of the cell.",
	"Water boils at 100 degrees Celsius at sea level."
	],
	"History": [
	"The Roman Empire fell in 476 AD.",
	"World War II ended in 1945.",
	"The Great Wall of China was built over centuries."
	],
	"Coding": [
	"Python uses indentation for blocks.",
	"React uses a virtual DOM for performance.",
	"Docker containers share the host OS kernel."
	]
	}

	facts = topics_data.get(TOPIC, ["Generic fact 1", "Generic fact 2"])

	for i, text in enumerate(facts):
	vec = embedder.encode(text)
	KNOWLEDGE_BASE.append({
	"id": f"{SHARD_ID}_{i}",
	"text": text,
	"vector": vec,
	"metadata": {"centrality": 0.9, "recency": 1.0, "weight": 1.0}
	})
	print(f"✅ Loaded {len(KNOWLEDGE_BASE)} facts.")

	@app.on_event("startup")
	async def startup_event():
	seed_knowledge()

	# --- API ---

	class RetrievalRequest(BaseModel):
	query_text: str
	query_vector: List[float] = None # Optional, if router computed it

	class RetrievalResponse(BaseModel):
	shard_id: str
	best_text: str
	score: float

	@app.post("/retrieve", response_model=RetrievalResponse)
	async def retrieve(req: RetrievalRequest):
	# 1. Get Query Vector
	if req.query_vector:
	q_vec = np.array(req.query_vector)
	else:
	q_vec = embedder.encode(req.query_text)

	# 2. Vector Search (Dot Product)
	# Stack all vectors
	db_vecs = np.array([item["vector"] for item in KNOWLEDGE_BASE])

	# Cosine Sim
	norm_q = np.linalg.norm(q_vec)
	norm_db = np.linalg.norm(db_vecs, axis=1)
	dot = np.dot(db_vecs, q_vec)
	sims = dot / (norm_q * norm_db + 1e-9)

	# 3. Apply Metadata (P = S * C * R * W)
	best_score = -1.0
	best_idx = 0

	for i, sim in enumerate(sims):
	meta = KNOWLEDGE_BASE[i]["metadata"]
	p_score = sim * meta["centrality"] * meta["recency"] * meta["weight"]

	if p_score > best_score:
	best_score = p_score
	best_idx = i

	return RetrievalResponse(
	shard_id=SHARD_ID,
	best_text=KNOWLEDGE_BASE[best_idx]["text"],
	score=float(best_score)
	)

	if __name__ == "__main__":
	port = int(os.getenv("PORT", 8001))
	uvicorn.run(app, host="0.0.0.0", port=port)