fix model

app.py

@@ -6,6 +6,7 @@ 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"})
 
@@ -18,25 +19,30 @@ def clean_and_extract_json(text):
     text = re.sub(r"```json|```", "", text).strip()
     return 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(cleaned_text)
+        # Check fields and strip spaces
         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")
+            raise ValueError("Missing or invalid fields in JSON")
         return {
             "script": script,
             "caption": caption,
             "hashtags": hashtags
         }
     except Exception:
+        # Fallback: don't include the prompt, just first sentence or default text
         fallback_script = raw_text.strip()
+        # Remove prompt repetition if possible
+        if prompt in fallback_script:
+            fallback_script = fallback_script.replace(prompt, "").strip()
         fallback_caption = fallback_script.split(".")[0].strip() or "Simak info menarik ini!"
         return {
             "script": fallback_script,
@@ -44,23 +50,13 @@ def generate_ui(prompt):
             "hashtags": ["#indonesia", "#fyp"]
         }
 
-# 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()
+demo = gr.Interface(
+    fn=generate_ui,
+    inputs=gr.Textbox(lines=6, placeholder="Masukkan prompt..."),
+    outputs=gr.JSON(),
+    title="TikTok Script Generator",
+    description="Generates short Indonesian scripts (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", port=7860, enable_api=True) #add api
+    demo.launch(server_name="0.0.0.0", port=7860, enable_api=True)