from fastapi import FastAPI
from pydantic import BaseModel  # Import BaseModel untuk mendefinisikan model data
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

app = FastAPI()

# Load model and tokenizer tanpa mendefinisikan cache_dir
tokenizer = AutoTokenizer.from_pretrained("cahya/gpt2-small-indonesian-522M")
model = AutoModelForCausalLM.from_pretrained("cahya/gpt2-small-indonesian-522M").to("cpu")

# Definisikan model data untuk body JSON
class GenerateRequest(BaseModel):
    prompt: str

@app.get("/")
def home():
    return {"message": "FastAPI running with Llama-3.2-1B-Instruct"}

@app.post("/generate")
def generate_text(request: GenerateRequest):  # Gunakan model data sebagai parameter
    inputs = tokenizer(request.prompt, return_tensors="pt").to("cpu")  # Ambil prompt dari request
    output = model.generate(**inputs, max_length=300)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return {"generated_text": generated_text}