Update script_gen.py

script_gen.py

@@ -58,7 +58,10 @@ def _get_client() -> InferenceClient:
             "HF_TOKEN environment variable is not set. "
             "Please set your Hugging Face API token to use the script generation feature."
         )
-    return InferenceClient(token=token)
+    return InferenceClient(
+        provider="hf-inference",
+        token=token,
+    )
 
 
 def generate_script(
@@ -95,28 +98,22 @@ def generate_script(
 
     client = _get_client()
 
-    # Call the model using text_generation
-    response = client.text_generation(
-        prompt=_format_prompt(SYSTEM_PROMPT, user_message),
+    # Call the model using chat_completion
+    response = client.chat_completion(
         model=MODEL_ID,
-        max_new_tokens=MAX_NEW_TOKENS,
+        messages=[
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": user_message},
+        ],
+        max_tokens=MAX_NEW_TOKENS,
         temperature=TEMPERATURE,
         top_p=0.9,
-        do_sample=True,
     )
 
-    script = response.strip()
+    script = response.choices[0].message.content.strip()
 
     if not script:
         raise RuntimeError("The model returned an empty script. Please try again.")
 
     logger.info("Script generated: %d characters", len(script))
-    return script
-
-
-def _format_prompt(system: str, user: str) -> str:
-    """
-    Format a prompt for SmolLM3 instruction following.
-    Uses a simple system + user format.
-    """
-    return f"<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n"
+    return script