Update app.py

app.py

@@ -1,54 +1,208 @@
-from fastapi import FastAPI, Request
+# from fastapi import FastAPI, Request
+# from pydantic import BaseModel
+# from transformers import AutoModel, AutoTokenizer
+# import torch
+
+# app = FastAPI()
+
+# model_id = "Qwen/Qwen3-Embedding-0.6B"
+
+# # Load tokenizer
+# tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# # Load model with GPU if available, else CPU
+# use_gpu = torch.cuda.is_available()
+
+# if use_gpu:
+#     print("CUDA is available, loading model with 4-bit quantization on GPU.")
+#     model = AutoModel.from_pretrained(
+#         model_id,
+#         device_map="auto",
+#         torch_dtype=torch.float16,
+#         load_in_4bit=True
+#     )
+# else:
+#     print("CUDA not available, loading model without 4-bit quantization on CPU.")
+#     model = AutoModel.from_pretrained(
+#         model_id,
+#         device_map="cpu",
+#         torch_dtype=torch.float32
+#     )
+
+# model.eval()
+
+# class TextInput(BaseModel):
+#     text: str
+
+# @app.post("/embed")
+# async def embed_text(input: TextInput):
+#     inputs = tokenizer(input.text, return_tensors="pt", truncation=True, max_length=512)
+
+#     # Move input tensors to same device as model
+#     device = next(model.parameters()).device
+#     inputs = {k: v.to(device) for k, v in inputs.items()}
+
+#     with torch.no_grad():
+#         outputs = model(**inputs)
+#         embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
+
+#     # Convert to list for JSON serialization
+#     return {"embedding": embeddings[0].cpu().tolist()}
+
+# from fastapi import FastAPI
+# from pydantic import BaseModel
+# from typing import List
+# from transformers import AutoTokenizer, AutoModel
+# import torch
+# import torch.nn.functional as F
+
+# app = FastAPI()
+
+# # Model config
+# MODEL_ID = "Qwen/Qwen3-Embedding-0.6B"
+# USE_GPU = torch.cuda.is_available()
+
+# # Load tokenizer
+# tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side='left')
+
+# # Load model with appropriate settings
+# if USE_GPU:
+#     print("🔋 Loading model on GPU with 4-bit quantization...")
+#     model = AutoModel.from_pretrained(
+#         MODEL_ID,
+#         device_map="auto",
+#         torch_dtype=torch.float16,
+#         load_in_4bit=True
+#     )
+# else:
+#     print("🧠 Loading model on CPU...")
+#     model = AutoModel.from_pretrained(
+#         MODEL_ID,
+#         device_map="cpu",
+#         torch_dtype=torch.float32
+#     )
+
+# model.eval()
+# device = next(model.parameters()).device
+
+# # Input schema
+# class EmbedRequest(BaseModel):
+#     texts: List[str]
+
+# # Output schema
+# class EmbedResponse(BaseModel):
+#     embeddings: List[List[float]]
+
+# # Masked mean pooling (ignores padded tokens)
+# def masked_mean_pooling(last_hidden_state, attention_mask):
+#     mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
+#     masked_embeddings = last_hidden_state * mask
+#     summed = masked_embeddings.sum(dim=1)
+#     counts = mask.sum(dim=1)
+#     return summed / counts.clamp(min=1e-9)
+
+# @app.post("/embed", response_model=EmbedResponse)
+# async def embed_texts(request: EmbedRequest):
+#     # Tokenize input texts
+#     inputs = tokenizer(
+#         request.texts,
+#         return_tensors="pt",
+#         padding=True,
+#         truncation=True,
+#         max_length=32768  # Qwen supports long sequences
+#     )
+#     inputs = {k: v.to(device) for k, v in inputs.items()}
+
+#     # Get embeddings
+#     with torch.no_grad():
+#         outputs = model(**inputs)
+#         pooled = masked_mean_pooling(outputs.last_hidden_state, inputs['attention_mask'])
+#         normalized = F.normalize(pooled, p=2, dim=1)
+
+#     return {"embeddings": normalized.cpu().tolist()}
+
+from fastapi import FastAPI, HTTPException, Header, Depends
 from pydantic import BaseModel
-from transformers import AutoModel, AutoTokenizer
+from typing import List, Optional
+from transformers import AutoTokenizer, AutoModel
 import torch
+import torch.nn.functional as F
+import os
+from dotenv import load_dotenv
 
-app = FastAPI()
+# Load environment variables from .env (if present)
+load_dotenv()
+
+# Load API key from environment variable
+API_KEY = os.environ.get("API_KEY")
+if not API_KEY:
+    raise RuntimeError("❌ EMBEDDING_API_KEY not set in environment!")
 
-model_id = "Qwen/Qwen3-Embedding-0.6B"
+# Initialize FastAPI
+app = FastAPI()
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+# Load tokenizer & model
+MODEL_ID = "Qwen/Qwen3-Embedding-0.6B"
+USE_GPU = torch.cuda.is_available()
 
-# Load model with GPU if available, else CPU
-use_gpu = torch.cuda.is_available()
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side='left')
 
-if use_gpu:
-    print("CUDA is available, loading model with 4-bit quantization on GPU.")
+if USE_GPU:
+    print("🔋 Using GPU with 4-bit quantization")
     model = AutoModel.from_pretrained(
-        model_id,
+        MODEL_ID,
         device_map="auto",
         torch_dtype=torch.float16,
         load_in_4bit=True
     )
 else:
-    print("CUDA not available, loading model without 4-bit quantization on CPU.")
+    print("🧠 Using CPU")
     model = AutoModel.from_pretrained(
-        model_id,
+        MODEL_ID,
         device_map="cpu",
         torch_dtype=torch.float32
     )
 
 model.eval()
+device = next(model.parameters()).device
+
+# Schema
+class EmbedRequest(BaseModel):
+    texts: List[str]
 
-class TextInput(BaseModel):
-    text: str
+class EmbedResponse(BaseModel):
+    embeddings: List[List[float]]
 
-@app.post("/embed")
-async def embed_text(input: TextInput):
-    inputs = tokenizer(input.text, return_tensors="pt", truncation=True, max_length=512)
+# Auth dependency
+async def verify_api_key(x_api_key: Optional[str] = Header(None)):
+    if x_api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
 
-    # Move input tensors to same device as model
-    device = next(model.parameters()).device
+# Masked mean pooling
+def masked_mean_pooling(last_hidden_state, attention_mask):
+    mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
+    masked_embeddings = last_hidden_state * mask
+    summed = masked_embeddings.sum(dim=1)
+    counts = mask.sum(dim=1)
+    return summed / counts.clamp(min=1e-9)
+
+# Endpoint
+@app.post("/embed", response_model=EmbedResponse, dependencies=[Depends(verify_api_key)])
+async def embed_texts(request: EmbedRequest):
+    inputs = tokenizer(
+        request.texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=32768
+    )
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
     with torch.no_grad():
         outputs = model(**inputs)
-        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
-
-    # Convert to list for JSON serialization
-    return {"embedding": embeddings[0].cpu().tolist()}
-
+        pooled = masked_mean_pooling(outputs.last_hidden_state, inputs['attention_mask'])
+        normalized = F.normalize(pooled, p=2, dim=1)
 
+    return {"embeddings": normalized.cpu().tolist()}