Axzyl commited on
Commit
6de89fc
·
verified ·
1 Parent(s): 650149e

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -6
app.py CHANGED
@@ -1,12 +1,8 @@
1
  from fastapi import FastAPI
2
  from pydantic import BaseModel
3
- from transformers import pipeline
4
 
5
  app = FastAPI(title="SAM_MedTesting")
6
 
7
- # Load your model once at startup
8
- generator = pipeline("text-generation", model="gpt2")
9
-
10
  class GenerationRequest(BaseModel):
11
  prompt: str
12
  max_new_tokens: int = 50
@@ -16,8 +12,8 @@ class GenerationResponse(BaseModel):
16
 
17
  @app.post("/generate", response_model=GenerationResponse)
18
  def generate(req: GenerationRequest):
19
- out = generator(req.prompt, max_length=req.max_new_tokens, do_sample=True)
20
- return GenerationResponse(generated_text=out[0]["generated_text"])
21
 
22
  @app.get("/health")
23
  def health():
 
1
  from fastapi import FastAPI
2
  from pydantic import BaseModel
 
3
 
4
  app = FastAPI(title="SAM_MedTesting")
5
 
 
 
 
6
  class GenerationRequest(BaseModel):
7
  prompt: str
8
  max_new_tokens: int = 50
 
12
 
13
  @app.post("/generate", response_model=GenerationResponse)
14
  def generate(req: GenerationRequest):
15
+ out = f"hello world: {req}"
16
+ return out
17
 
18
  @app.get("/health")
19
  def health():