Falln87 commited on
Commit
fd3f98d
·
verified ·
1 Parent(s): 186e06e

Create app.pu

Browse files
Files changed (1) hide show
  1. app.pu +14 -0
app.pu ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
3
+
4
+ app = FastAPI()
5
+
6
+ model_id = "gpt2" # Replace with your desired model
7
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
8
+ model = AutoModelForCausalLM.from_pretrained(model_id)
9
+ generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
10
+
11
+ @app.get("/")
12
+ def generate_text(prompt: str):
13
+ result = generator(prompt)
14
+ return {"generated_text": result[0]["generated_text"]}