cordia-api / app.py
petertill's picture
Update app.py
e18a7f6 verified
raw
history blame
1.06 kB
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import uvicorn
app = FastAPI()
try:
model = AutoModelForCausalLM.from_pretrained("petertill/cordia-a6")
tokenizer = AutoTokenizer.from_pretrained("petertill/cordia-a6")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
print("Model and tokenizer loaded successfully!")
class GenerateRequest(BaseModel):
prompt: str
class GenerateResponse(BaseModel):
generated_text: str
@app.post("/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
try:
output = pipe(request.prompt, max_length=200)[0]['generated_text']
return GenerateResponse(generated_text=output)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
except Exception as e:
print(f"Error: {e}")
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)