File size: 1,555 Bytes
bbd6b6a 6993d08 bbd6b6a 6993d08 bbd6b6a 6993d08 bbd6b6a 6993d08 bbd6b6a 6993d08 bbd6b6a 6993d08 bbd6b6a 6993d08 bbd6b6a 6993d08 bbd6b6a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
app = FastAPI()
# Switching to Gemma-2-2B-Instruct (High Quality & Good Speed)
MODEL_REPO = "bartowski/gemma-2-2b-it-GGUF"
MODEL_FILE = "gemma-2-2b-it-Q4_K_M.gguf"
print("Downloading Gemma-2-2B model...")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
print("Loading Gemma for high-quality CPU inference...")
llm = Llama(
model_path=model_path,
n_ctx=2048,
n_threads=2,
verbose=False
)
class PromptRequest(BaseModel):
prompt: str
@app.get("/")
def read_root():
return {"message": "Gemma-2-2B-IT API is running."}
@app.get("/health")
def health_check():
return {"status": "alive"}
@app.post("/api")
async def generate_response(request: PromptRequest):
try:
# Gemma 2 Prompt Format
formatted_prompt = f"<start_of_turn>user\n{request.prompt}<end_of_turn>\n<start_of_turn>model\n"
output = llm(
formatted_prompt,
max_tokens=1024,
stop=["<end_of_turn>"],
echo=False
)
response_text = output['choices'][0]['text'].strip()
return {
"status": "success",
"text": response_text
}
except Exception as e:
print(f"Error: {e}")
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|