UnityAPI / main.py
bp7274's picture
Update main.py
d354551 verified
raw
history blame contribute delete
972 Bytes
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
# Load model
generator = pipeline("text-generation", model="bp7274/soulprint-model")
app = FastAPI()
class InputData(BaseModel):
age: int
location: str
profession: str
personality: str
@app.get("/")
def health():
return {"status": "ready"}
@app.post("/generate")
def generate_quote(data: InputData):
try:
prompt = f"{data.age}, {data.location}, {data.profession}, {data.personality}."
result = generator(
prompt,
max_length=120,
temperature=0.9,
top_p=0.95,
repetition_penalty=1.2,
do_sample=True
)[0]["generated_text"]
# Strip only the prompt portion (no fallback check)
quote = result[len(prompt):].strip()
return {"quote": quote}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))