change model

app.py

@@ -2,41 +2,65 @@ import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import json
+import re
 
 MODEL_ID = "cahya/gpt2-small-indonesian-522M"
 
-# Load tokenizer + model on CPU
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32, device_map={"": "cpu"})
 
-generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer
-)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-def generate(prompt):
-    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"]
+def clean_and_extract_json(text):
+    json_start = text.find('{')
+    if json_start >= 0:
+        text = text[json_start:]
+    text = re.sub(r"```json|```", "", text).strip()
+    return text
 
-    # Try to parse the output as JSON, fallback to raw text
+# This function is used in the UI
+def generate_ui(prompt):
+    out = generator(prompt, max_new_tokens=180, do_sample=True, top_p=0.95, top_k=40, temperature=0.8)
+    raw_text = out[0]["generated_text"]
+    cleaned_text = clean_and_extract_json(raw_text)
     try:
-        parsed = json.loads(text)
-        return parsed
+        parsed = json.loads(cleaned_text)
+        script = parsed.get("script", "").strip()
+        caption = parsed.get("caption", "").strip()
+        hashtags = parsed.get("hashtags", [])
+        if not script or not caption or not isinstance(hashtags, list):
+            raise ValueError("Missing fields")
+        return {
+            "script": script,
+            "caption": caption,
+            "hashtags": hashtags
+        }
     except Exception:
+        fallback_script = raw_text.strip()
+        fallback_caption = fallback_script.split(".")[0].strip() or "Simak info menarik ini!"
         return {
-            "script": text.strip(),
-            "caption": "Generated caption",
+            "script": fallback_script,
+            "caption": fallback_caption,
             "hashtags": ["#indonesia", "#fyp"]
         }
 
-demo = gr.Interface(
-    fn=generate,
-    inputs=gr.Textbox(lines=6, placeholder="Prompt..."),
-    outputs=gr.JSON(),
-    title="TikTok Script Generator",
-    description="Generates short Indonesian scripts (JSON)."
+# This is the API function Gradio will call, input/output signature must match components exactly
+@gr.api(
+    input=gr.Textbox(lines=6),
+    output=gr.JSON()
 )
+def generate_api(prompt):
+    return generate_ui(prompt)
+
+with gr.Blocks() as demo:
+    prompt_input = gr.Textbox(lines=6, placeholder="Prompt...")
+    output_json = gr.JSON()
+
+    generate_btn = gr.Button("Generate")
+    generate_btn.click(fn=generate_ui, inputs=prompt_input, outputs=output_json)
+
+    gr.Markdown("### TikTok Script Generator")
+    gr.Markdown("Generates short Indonesian scripts (JSON).")
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.launch(server_name="0.0.0.0", port=7860, enable_api=True)