| 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) |
