Update app.py

app.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI
+from pydantic import BaseModel  # Import BaseModel untuk mendefinisikan model data
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
@@ -8,13 +9,17 @@ app = FastAPI()
 tokenizer = AutoTokenizer.from_pretrained("unsloth/Llama-3.2-1B-Instruct")
 model = AutoModelForCausalLM.from_pretrained("unsloth/Llama-3.2-1B-Instruct").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(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+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}
+    return {"generated_text": generated_text}