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)