Merge 8607169546 into dddedc6881

Automated security fix generated by Orbis Security AI

model/model_minimind.py

@@ -266,7 +266,8 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
             attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
             logits = outputs.logits[:, -1, :] / temperature
             if repetition_penalty != 1.0:
-                for i in range(input_ids.shape[0]): logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+                for i in range(input_ids.shape[0]):
+                    seen = torch.unique(input_ids[i]); score = logits[i, seen]; logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
             if top_k > 0: 
                 logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
             if top_p < 1.0:

scripts/serve_openai_api.py

@@ -13,8 +13,9 @@ import uvicorn
 
 from threading import Thread
 from queue import Queue
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Depends
 from fastapi.responses import StreamingResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
@@ -24,6 +25,14 @@ warnings.filterwarnings('ignore')
 
 app = FastAPI()
 
+_security = HTTPBearer(auto_error=False)
+
+
+def _verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(_security)):
+    api_key = os.environ.get("MINIMIND_API_KEY")
+    if api_key and (credentials is None or credentials.credentials != api_key):
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+
 
 def init_model(args):
     tokenizer = AutoTokenizer.from_pretrained(args.load_from)
@@ -169,7 +178,7 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
 
 
 @app.post("/v1/chat/completions")
-async def chat_completions(request: ChatRequest):
+async def chat_completions(request: ChatRequest, _: None = Depends(_verify_api_key)):
     try:
         if request.stream:
             return StreamingResponse(