api-gpt / app.py
fikrialfanf's picture
Update app.py
dcaa07a verified
Raw
History Blame Contribute Delete
2.5 kB
import os
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import pipeline, set_seed
# Direktori cache Hugging Face
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
app = FastAPI(title="StemdjaGPT", version="1.0")
# Middleware CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
MODEL = "cahya/gpt2-small-indonesian-522M"
generator = pipeline("text-generation", model=MODEL, tokenizer=MODEL)
set_seed(42)
class PromptInput(BaseModel):
prompt: str
max_tokens: int = 256 # ganti nama parameter agar lebih jelas
temperature: float = 0.8
@app.post("/generate")
def generate_text(item: PromptInput):
result = generator(
item.prompt,
max_new_tokens=item.max_tokens, # pakai max_new_tokens, bukan max_length
do_sample=True,
temperature=item.temperature,
top_p=0.9,
truncation=True
)[0]["generated_text"]
return {
"input": item.prompt,
"output": result.strip()
}
@app.get("/")
def home():
return {"message": "StemdjaGPT siap menghasilkan teks Bahasa Indonesia"}
import os
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import pipeline, set_seed
# Direktori cache Hugging Face
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
app = FastAPI(title="StemdjaGPT", version="1.0")
# Middleware CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
MODEL = "cahya/gpt2-small-indonesian-522M"
generator = pipeline("text-generation", model=MODEL, tokenizer=MODEL)
set_seed(42)
class PromptInput(BaseModel):
prompt: str
max_tokens: int = 256 # ganti nama parameter agar lebih jelas
temperature: float = 0.8
@app.post("/generate")
def generate_text(item: PromptInput):
result = generator(
item.prompt,
max_new_tokens=item.max_tokens, # pakai max_new_tokens, bukan max_length
do_sample=True,
temperature=item.temperature,
top_p=0.9,
truncation=True
)[0]["generated_text"]
return {
"input": item.prompt,
"output": result.strip()
}
@app.get("/")
def home():
return {"message": "StemdjaGPT siap menghasilkan teks Bahasa Indonesia"}