| |
| from fastapi import FastAPI, HTTPException |
| from fastapi.responses import FileResponse |
| from fastapi.middleware.cors import CORSMiddleware |
| from pydantic import BaseModel |
| import subprocess |
| import tiktoken |
| import os |
| import time |
|
|
| app = FastAPI() |
|
|
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
| class GenerateRequest(BaseModel): |
| prompt: str |
| max_tokens: int = 100 |
| temperature: float = 0.8 |
| top_k: int = 40 |
|
|
| try: |
| enc = tiktoken.get_encoding("gpt2") |
| print("Tokenizer loaded successfully.") |
| except Exception as e: |
| print(f"Warning: tiktoken not found. Error: {e}") |
| enc = None |
|
|
| @app.get("/") |
| async def root(): |
| current_dir = os.path.dirname(os.path.abspath(__file__)) |
| return FileResponse(os.path.join(current_dir, "index.html")) |
|
|
| @app.get("/health") |
| async def health_check(): |
| current_dir = os.path.dirname(os.path.abspath(__file__)) |
| exe_path = os.path.join(current_dir, "inference") |
| model_path = os.path.join(current_dir, "model.bin") |
| return { |
| "status": "ok", |
| "inference_exe_found": os.path.exists(exe_path), |
| "model_bin_found": os.path.exists(model_path), |
| "working_directory": current_dir |
| } |
|
|
| @app.post("/generate") |
| async def generate_text(req: GenerateRequest): |
| if enc is None: |
| raise HTTPException(status_code=500, detail="Tokenizer not loaded.") |
|
|
| input_tokens = enc.encode(req.prompt) |
| token_str = ",".join(map(str, input_tokens)) |
|
|
| current_dir = os.path.dirname(os.path.abspath(__file__)) |
| exe_path = os.path.join(current_dir, "inference") |
| model_path = os.path.join(current_dir, "model.bin") |
|
|
| if not os.path.exists(exe_path): |
| raise HTTPException(status_code=500, detail=f"inference binary not found: {exe_path}") |
|
|
| if not os.path.exists(model_path): |
| raise HTTPException(status_code=500, detail=f"model.bin not found: {model_path}") |
|
|
| try: |
| start_time = time.perf_counter() |
| process = subprocess.run( |
| [exe_path, token_str, str(req.max_tokens), str(req.temperature), str(req.top_k)], |
| capture_output=True, |
| text=True, |
| cwd=current_dir |
| ) |
| elapsed_ms = (time.perf_counter() - start_time) * 1000 |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Execution failed: {str(e)}") |
|
|
| if process.returncode != 0 and not process.stdout.strip(): |
| stdout_msg = process.stdout.strip() if process.stdout else "" |
| stderr_msg = process.stderr.strip() if process.stderr else "" |
| raise HTTPException(status_code=500, detail=f"C++ Error | stdout: '{stdout_msg}' | stderr: '{stderr_msg}'") |
|
|
| try: |
| output_str = process.stdout.strip() |
| generated_ids = [] |
| if output_str: |
| for x in output_str.split(): |
| try: |
| generated_ids.append(int(x)) |
| except ValueError: |
| pass |
|
|
| generated_text = enc.decode(generated_ids) if generated_ids else "" |
| tokens_out = len(generated_ids) |
| tokens_per_sec = round(tokens_out / (elapsed_ms / 1000), 2) if elapsed_ms > 0 else 0 |
|
|
| return { |
| "prompt": req.prompt, |
| "generated_text": generated_text, |
| "tokens_in": len(input_tokens), |
| "tokens_out": tokens_out, |
| "latency_ms": round(elapsed_ms, 2), |
| "tokens_per_sec": tokens_per_sec |
| } |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=f"Decoding error: {str(e)}") |
