return json

app.py

@@ -1,11 +1,11 @@
-# app.py
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
+import json
 
-MODEL_ID = "cahya/gpt2-small-indonesian-522M"  # change if needed
+MODEL_ID = "cahya/gpt2-small-indonesian-522M"
 
-# Load tokenizer + model (CPU)
+# Load tokenizer + model on CPU
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32, device_map={"": "cpu"})
 
@@ -16,16 +16,24 @@ generator = pipeline(
 )
 
 def generate(prompt):
-    # prompt already instructs JSON-only output; but we keep generation stable
     out = generator(prompt, max_new_tokens=180, do_sample=True, top_p=0.95, top_k=40, temperature=0.8)
     text = out[0]["generated_text"]
-    return text
 
-# Simple gradio UI (the same function is exposed as API)
+    # Try to parse the output as JSON, fallback to raw text
+    try:
+        parsed = json.loads(text)
+        return parsed
+    except Exception:
+        return {
+            "script": text.strip(),
+            "caption": "Generated caption",
+            "hashtags": ["#indonesia", "#fyp"]
+        }
+
 demo = gr.Interface(
     fn=generate,
     inputs=gr.Textbox(lines=6, placeholder="Prompt..."),
-    outputs=gr.Textbox(),
+    outputs=gr.JSON(),  # JSON output for proper API response
     title="TikTok Script Generator",
     description="Generates short Indonesian scripts (JSON)."
 )