Update main.py

main.py

@@ -14,56 +14,57 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 import torch
 
 # ==========================================
-# OPTIMASI CORE (2 vCPU / 16GB RAM)
+# KONFIGURASI & OPTIMASI CPU
 # ==========================================
 torch.set_num_threads(1)
 inference_pool = ThreadPoolExecutor(max_workers=2)
 ml_models = {}
 
-# Konfigurasi Identitas
-MODEL_ID = "trico-a-522m"
-MODEL_NAME = "Trico A 522m"
+# Rebranding Info
+REAL_MODEL_ID = "abdiharyadi/cahya-gpt2-large-indonesian-522M-instruct"
+DISPLAY_MODEL_ID = "trico-a-522m"
 POWERED_BY = "SFTNexus"
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    print(f"🚀 [STARTUP] Memuat {MODEL_NAME} by {POWERED_BY}...")
-    
-    # Path model disesuaikan dengan input Anda
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    print(f"🚀 [STARTUP] Memuat {DISPLAY_MODEL_ID} (Powered by {POWERED_BY})...")
+    tokenizer = AutoTokenizer.from_pretrained(REAL_MODEL_ID)
     with torch.no_grad():
         model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID, 
+            REAL_MODEL_ID, 
             torch_dtype=torch.float32,
             low_cpu_mem_usage=True
         ).eval()
-    
     ml_models["tokenizer"] = tokenizer
     ml_models["model"] = model
-    print(f"✅ [STARTUP] {MODEL_NAME} siap di Hugging Face Spaces!")
     yield
     ml_models.clear()
     inference_pool.shutdown(wait=True)
 
 # ==========================================
-# AUTHENTIKASI & SCHEMA OPENAI
+# AUTHENTICATION
 # ==========================================
 VALID_API_KEYS = set(os.getenv("APIKEY", "sk-default").split(","))
 api_key_header = APIKeyHeader(name="Authorization", auto_error=True)
 
 def verify_api_key(auth_header: str = Security(api_key_header)):
     # Mendukung format "Bearer sk-..."
-    token = auth_header.replace("Bearer ", "") if "Bearer " in auth_header else auth_header
+    token = auth_header.replace("Bearer ", "")
     if token not in VALID_API_KEYS:
         raise HTTPException(status_code=401, detail="Invalid API Key")
     return token
 
+app = FastAPI(title=f"Trico A API", lifespan=lifespan)
+
+# ==========================================
+# OPENAI COMPATIBLE SCHEMAS
+# ==========================================
 class ChatMessage(BaseModel):
     role: str
     content: str
 
 class ChatCompletionRequest(BaseModel):
-    model: str = MODEL_ID
+    model: str = DISPLAY_MODEL_ID
     messages: List[ChatMessage]
     temperature: Optional[float] = 0.7
     top_p: Optional[float] = 0.9
@@ -71,117 +72,112 @@ class ChatCompletionRequest(BaseModel):
     stream: Optional[bool] = False
 
 # ==========================================
-# UTILS & PROMPT ENGINEERING
+# PROMPT ENGINEERING (CHAT TO ALPACA)
 # ==========================================
-def format_chat_to_alpaca(messages: List[ChatMessage]) -> str:
-    """Mengonversi chat history ke format instruksi agar model tidak 'ngaco'."""
-    prompt = f"System: Anda adalah {MODEL_NAME}, asisten AI native Indonesia yang handal. Powered by {POWERED_BY}.\n\n"
-    for msg in messages:
-        if msg.role == "user":
-            prompt += f"### Instruksi:\n{msg.content}\n\n"
-        elif msg.role == "assistant":
-            prompt += f"### Respons:\n{msg.content}\n\n"
-    prompt += "### Respons:\n"
-    return prompt
+def format_openai_to_alpaca(messages: List[ChatMessage]) -> str:
+    # Mengambil pesan terakhir sebagai instruksi, sisanya sebagai context jika ada
+    instruction = messages[-1].content
+    return (
+        "Di bawah ini adalah instruksi yang menjelaskan suatu tugas. "
+        "Tuliskan respons yang melengkapi permintaan tersebut dengan tepat.\n\n"
+        f"### Instruksi:\n{instruction}\n\n### Respons:\n"
+    )
 
-async def openai_stream_generator(request_id: str, prompt: str, params: dict):
+# ==========================================
+# CORE INFERENCE (STREAMING & NON-STREAMING)
+# ==========================================
+async def chat_generator(request: ChatCompletionRequest):
     tokenizer = ml_models["tokenizer"]
     model = ml_models["model"]
+    created_time = int(time.time())
+    completion_id = f"chatcmpl-{uuid.uuid4()}"
     
+    prompt = format_openai_to_alpaca(request.messages)
     inputs = tokenizer([prompt], return_tensors="pt")
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
     gen_kwargs = dict(
         **inputs,
         streamer=streamer,
-        max_new_tokens=params["max_tokens"],
-        temperature=params["temperature"],
-        top_p=params["top_p"],
-        do_sample=params["temperature"] > 0,
+        max_new_tokens=request.max_tokens,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        do_sample=request.temperature > 0,
         pad_token_id=tokenizer.eos_token_id
     )
     
     loop = asyncio.get_running_loop()
     loop.run_in_executor(inference_pool, lambda: model.generate(**gen_kwargs))
+
+    for token in streamer:
+        chunk = {
+            "id": completion_id,
+            "object": "chat.completion.chunk",
+            "created": created_time,
+            "model": DISPLAY_MODEL_ID,
+            "system_fingerprint": f"powered_by_{POWERED_BY}",
+            "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}]
+        }
+        yield f"data: {json.dumps(chunk)}\n\n"
+        await asyncio.sleep(0.01)
     
-    created_time = int(time.time())
-    for chunk in streamer:
-        if chunk:
-            data = {
-                "id": request_id,
-                "object": "chat.completion.chunk",
-                "created": created_time,
-                "model": MODEL_ID,
-                "choices": [{"index": 0, "delta": {"content": chunk}, "finish_reason": None}]
-            }
-            yield f"data: {json.dumps(data)}\n\n"
-            await asyncio.sleep(0.01)
-            
     yield "data: [DONE]\n\n"
 
 # ==========================================
 # ENDPOINTS
 # ==========================================
-app = FastAPI(title=f"{MODEL_NAME} API", lifespan=lifespan)
-
 @app.get("/health")
 async def health():
-    return {"status": "ready", "model": MODEL_NAME, "provider": POWERED_BY}
+    return {"status": "ok", "model": DISPLAY_MODEL_ID, "provider": POWERED_BY}
+
+@app.get("/v1/models")
+async def list_models():
+    return {
+        "object": "list",
+        "data": [{
+            "id": DISPLAY_MODEL_ID,
+            "object": "model",
+            "created": 1700000000,
+            "owned_by": POWERED_BY
+        }]
+    }
 
 @app.post("/v1/chat/completions")
-@app.post("/generate") # Menjaga kompatibilitas request sebelumnya
-async def chat_completions(req: ChatCompletionRequest, api_key: str = Depends(verify_api_key)):
-    prompt = format_chat_to_alpaca(req.messages)
-    request_id = f"chatcmpl-{uuid.uuid4()}"
-    
-    # Mode Streaming (SSE)
-    if req.stream:
-        return StreamingResponse(
-            openai_stream_generator(request_id, prompt, req.model_dump()),
-            media_type="text/event-stream"
-        )
+async def chat_completions(request: ChatCompletionRequest, _ = Depends(verify_api_key)):
+    if request.stream:
+        return StreamingResponse(chat_generator(request), media_type="text/event-stream")
     
-    # Mode Non-Streaming
+    # Non-streaming Logic
     tokenizer = ml_models["tokenizer"]
     model = ml_models["model"]
-    
+    prompt = format_openai_to_alpaca(request.messages)
     inputs = tokenizer([prompt], return_tensors="pt")
-    loop = asyncio.get_running_loop()
     
-    output_tokens = await loop.run_in_executor(
+    loop = asyncio.get_running_loop()
+    output_ids = await loop.run_in_executor(
         inference_pool, 
         lambda: model.generate(
             **inputs, 
-            max_new_tokens=req.max_tokens,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            do_sample=req.temperature > 0,
+            max_new_tokens=request.max_tokens,
+            temperature=request.temperature,
+            do_sample=request.temperature > 0,
             pad_token_id=tokenizer.eos_token_id
         )
     )
     
-    full_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
-    # Ambil hanya bagian setelah Respons terakhir
+    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    # Ambil hanya bagian setelah "### Respons:"
     response_text = full_text.split("### Respons:\n")[-1].strip()
     
     return {
-        "id": request_id,
+        "id": f"chatcmpl-{uuid.uuid4()}",
         "object": "chat.completion",
         "created": int(time.time()),
-        "model": MODEL_ID,
+        "model": DISPLAY_MODEL_ID,
         "choices": [{
-            "index": 0,
             "message": {"role": "assistant", "content": response_text},
-            "finish_reason": "stop"
+            "finish_reason": "stop",
+            "index": 0
         }],
-        "usage": {
-            "prompt_tokens": len(inputs['input_ids'][0]),
-            "completion_tokens": len(output_tokens[0]) - len(inputs['input_ids'][0]),
-            "total_tokens": len(output_tokens[0])
-        },
-        "system_fingerprint": POWERED_BY
-    }
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} # Placeholder
+    }