mirror of
https://github.com/NousResearch/atropos.git
synced 2026-04-19 12:57:58 +00:00
28 lines
824 B
Python
28 lines
824 B
Python
from fastapi import FastAPI, Request
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
app = FastAPI()
|
|
|
|
# Load the model and tokenizer
|
|
model = AutoModelForCausalLM.from_pretrained("./opt-350m")
|
|
tokenizer = AutoTokenizer.from_pretrained("./opt-350m")
|
|
|
|
|
|
@app.post("/v1/chat/completions")
|
|
async def chat_completions(request: Request):
|
|
data = await request.json()
|
|
messages = data.get("messages", [])
|
|
# Extract prompt for potential future use
|
|
_ = messages[-1]["content"] if messages else ""
|
|
|
|
# Simulate a numerical rating (0-1) for now
|
|
response_text = "0.9"
|
|
|
|
# Return in OpenAI-compatible format
|
|
return {"choices": [{"message": {"content": response_text}}]}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="localhost", port=9001)
|