|
|
from fastapi import FastAPI |
|
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
from vllm import LLM, SamplingParams |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=["*"], |
|
|
allow_methods=["*"], |
|
|
allow_headers=["*"], |
|
|
) |
|
|
|
|
|
llama = LLM( |
|
|
model="meta-llama/Llama-3.2-1B-Instruct", |
|
|
dtype="half", |
|
|
gpu_memory_utilization=0.9 |
|
|
) |
|
|
|
|
|
@app.post("/generate") |
|
|
def generate(prompt: str): |
|
|
params = SamplingParams(temperature=0.7, max_tokens=100) |
|
|
output = llama.generate([prompt], params) |
|
|
return {"text": output[0].outputs[0].text} |
|
|
|
|
|
if __name__ == "__main__": |
|
|
import uvicorn |
|
|
uvicorn.run(app, host="0.0.0.0", port=7860) |
