Update utility/script_generator.py

utility/script_generator.py

@@ -1,50 +1,61 @@
-import g4f
-from g4f.client import Client
-import json
-
-def generate_script(topic):
-    prompt = (
-        """You are a seasoned content writer for a YouTube Shorts channel, specializing in facts videos. 
-        Your facts shorts are concise, each lasting less than 50 seconds (approximately 140 words). 
-        They are incredibly engaging and original. When a user requests a specific type of facts short, you will create it.
-
-        For instance, if the user asks for:
-        Weird facts
-        You would produce content like this:
-
-        Weird facts you don't know:
-        - Bananas are berries, but strawberries aren't.
-        - A single cloud can weigh over a million pounds.
-        - There's a species of jellyfish that is biologically immortal.
-        - Honey never spoils; archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still edible.
-        - The shortest war in history was between Britain and Zanzibar on August 27, 1896. Zanzibar surrendered after 38 minutes.
-        - Octopuses have three hearts and blue blood.
-
-        You are now tasked with creating the best short script based on the user's requested type of 'facts'.
-
-        Keep it brief, highly interesting, and unique.
-
-        Strictly output the script in a JSON format like below, and only provide a parsable JSON object with the key 'script'.
-
-        # Output
-        {"script": "Here is the script ..."}
-        """
-    )
-
-    client = Client()
-    response = client.chat.completions.create(
-        model='gpt-4o',
-        messages=[{'role': 'user', 'content': prompt + "\n\n" + topic}]
-        
-    )
-    
-    content = response.choices[0].message.content
-    try:
-        script = json.loads(content)["script"]
-    except json.JSONDecodeError:
-        print("JSONDecodeError. Attempting to extract JSON from the response.")
-        json_start = content.find('{')
-        json_end = content.rfind('}') + 1
-        script = json.loads(content[json_start:json_end])["script"]
-    
-    return script
+from transformers import pipeline
+import json
+
+def generate_script(topic):
+    prompt = (
+        """You are a seasoned content writer for a YouTube Shorts channel, specializing in facts videos. 
+        Your facts shorts are concise, each lasting less than 50 seconds (approximately 140 words). 
+        They are incredibly engaging and original. When a user requests a specific type of facts short, you will create it.
+
+        For instance, if the user asks for:
+        Weird facts
+        You would produce content like this:
+
+        Weird facts you don't know:
+        - Bananas are berries, but strawberries aren't.
+        - A single cloud can weigh over a million pounds.
+        - There's a species of jellyfish that is biologically immortal.
+        - Honey never spoils; archaeologists have found pots of honey in ancient Egyptian tombs that are over 3,000 years old and still edible.
+        - The shortest war in history was between Britain and Zanzibar on August 27, 1896. Zanzibar surrendered after 38 minutes.
+        - Octopuses have three hearts and blue blood.
+
+        You are now tasked with creating the best short script based on the user's requested type of 'facts'.
+
+        Keep it brief, highly interesting, and unique.
+
+        Strictly output the script in a JSON format like below, and only provide a parsable JSON object with the key 'script'.
+
+        # Output
+        {"script": "Here is the script ..."}
+        """
+    )
+
+    # Initialize the text-generation pipeline
+    pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B")
+
+    # Generate the script using the pipeline
+    response = pipe(
+        [{"role": "user", "content": prompt + "\n\n" + topic}],
+        max_length=500,  # Adjust the max_length as needed
+        num_return_sequences=1,
+        truncation=True
+    )
+
+    # Extract the generated text
+    generated_text = response[0]['generated_text']
+
+    try:
+        # Attempt to parse the generated text as JSON
+        script = json.loads(generated_text)["script"]
+    except json.JSONDecodeError:
+        print("JSONDecodeError. Attempting to extract JSON from the response.")
+        json_start = generated_text.find('{')
+        json_end = generated_text.rfind('}') + 1
+        script = json.loads(generated_text[json_start:json_end])["script"]
+    
+    return script
+
+# # Example usage
+# topic = "Weird facts"
+# script = generate_script(topic)
+# print(script)