|
from fastapi import FastAPI, Request |
|
from fastapi.responses import JSONResponse, FileResponse |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
import uvicorn |
|
|
|
app = FastAPI() |
|
|
|
|
|
model_name = "microsoft/DialoGPT-small" |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
model = AutoModelForCausalLM.from_pretrained(model_name) |
|
|
|
|
|
@app.get("/") |
|
async def serve_index(): |
|
return FileResponse("static/index.html") |
|
|
|
|
|
@app.post("/chat") |
|
async def chat(request: Request): |
|
data = await request.json() |
|
prompt = data.get("prompt", "") |
|
|
|
|
|
inputs = tokenizer(prompt, return_tensors="pt") |
|
outputs = model.generate(inputs["input_ids"], max_length=50) |
|
response = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
return JSONResponse({"response": response}) |
|
|
|
if __name__ == "__main__": |
|
uvicorn.run(app, host="0.0.0.0", port=7860) |
|
|