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"}