FlameF0X commited on
Commit
f8898c0
·
verified ·
1 Parent(s): 53e4d1a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -0
app.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, Request
2
+ from pydantic import BaseModel
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ import torch
5
+
6
+ app = FastAPI()
7
+
8
+ # Load model and tokenizer once at startup
9
+ model_name = "gpt2" # change this to your own model
10
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
11
+ model = AutoModelForCausalLM.from_pretrained(model_name)
12
+
13
+ class PromptRequest(BaseModel):
14
+ prompt: str
15
+ max_new_tokens: int = 50
16
+
17
+ @app.post("/generate")
18
+ async def generate_text(req: PromptRequest):
19
+ inputs = tokenizer(req.prompt, return_tensors="pt")
20
+ outputs = model.generate(
21
+ **inputs,
22
+ max_new_tokens=req.max_new_tokens,
23
+ do_sample=True,
24
+ temperature=0.8,
25
+ top_p=0.95,
26
+ )
27
+ generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
28
+ return {"generated_text": generated}