Spaces:
Sleeping
Sleeping
File size: 1,059 Bytes
e18a7f6 11b89be 78da8bb e18a7f6 78da8bb 11b89be e18a7f6 11b89be e18a7f6 11b89be e18a7f6 11b89be e18a7f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
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) |