Update recommendations_and_engagement_predict.py

recommendations_and_engagement_predict.py

@@ -1,127 +1,127 @@
-import os
-
-
-import gradio as gr
-# from transformers import AutoTokenizer, AutoModelForCausalLM
-from sentence_transformers import SentenceTransformer
-from datasets import load_dataset
-import numpy as np
-import torch
-from shared_resources import shared_resources
-from phi.agent import Agent
-from phi.tools.duckduckgo import DuckDuckGo
-from phi.agent import Agent, RunResponse
-from phi.model.huggingface import HuggingFaceChat
-
-
-
-class StoryRecommendationApp:
-
-    def __init__(self):
-        # self.tokenizer = shared_resources.tokenizer
-        # self.model = shared_resources.model
-        self.device = shared_resources.device
-        self.sentence_transformer = shared_resources.sentence_transformer
-        self.data = shared_resources.data
-
-    def search(self, query: str, k: int = 3):
-        """Search for recommended videos based on user input."""
-        embedded_query = self.sentence_transformer.encode(query)
-        scores, retrieved_examples = self.data.get_nearest_examples("embeddings", embedded_query, k=k)
-
-        if isinstance(retrieved_examples, np.ndarray):
-            retrieved_examples = np.nan_to_num(retrieved_examples, nan=0)
-        elif isinstance(retrieved_examples, list):
-            retrieved_examples = [
-                [0 if np.isnan(val) else val for val in example] if isinstance(example, list) else example
-                for example in retrieved_examples
-            ]
-        return scores, retrieved_examples
-
-    def compute_mean_and_predict(self, retrieved_examples):
-        """Compute mean for likesCount, commentCount, and shareCount."""
-        features = ["LikesCount", "commentCount", "shareCount"]
-        predictions = {}
-
-        for feature in features:
-            values = np.array(retrieved_examples[feature])
-            values = np.array([np.nan if v is None else v for v in values])
-            values = np.nan_to_num(values, nan=0.0)
-
-            mean_value = int(np.mean(values))
-            predictions[f"predicted_{feature}"] = mean_value
-
-        return predictions
-
-    def generate_prompt(self, query: str):
-        """Generate a prompt for video generation based on the input story."""
-        input_text = f'''
-        I want to summarize a story in exactly 3 sentences. No more than 3 nor less than 3. 
-        But the sentences have to be good enough to use as a prompt for video generation, because I have to give those 3 sentences to the video generation model.
-        For example: This prompt is about A heartwarming family reunion celebrating love and cherished memories in exactly 3 sentences.
-        -A warm, heartfelt reunion in a cozy living room, where family members embrace each other after a long time apart, soft lighting enhances the intimate atmosphere, and laughter fills the air.
-        -A close-up shot of a grandmother’s hands carefully arranging a family photo album, as the camera pans over old pictures, evoking cherished memories and a deep sense of love.
-        -A final moment around the dinner table, family members sharing a meal together, toasts are made, and the soft glow of candles reflects the joy and connection between generations.
-
-        So, I will provide you that story now. The story is:\n{query}
-
-        Please remember, don't give any background descriptions. Just generate 3 sentences likewise the example above. Don't even give the starting text like: "Here are the 3 sentences that summarize the story:" or any other like this. Just give the answers in 3 sentences directly
-        '''
-        agent = Agent(
-            model=HuggingFaceChat(
-                id="meta-llama/Meta-Llama-3-8B-Instruct",
-                max_tokens=4096,
-            ),
-            # tools=[DuckDuckGo()],
-            markdown=True
-        )
-
-        # Get the response in a variable
-        run: RunResponse = agent.run(input_text)
-        generated_text=run.content
-
-        sentences = [sentence.strip() for sentence in generated_text.split('.') if sentence]
-        return '. '.join(sentences[-3:]) + ('.' if len(sentences) > 0 else '')
-
-    def generate_story_and_recommendation(self, generated_response: str):
-        """Generate story recommendations and predictions based on the user input."""
-        scores, result = self.search(generated_response, 4)
-
-        recommended_videos_text = ""
-        predictions = {}
-
-        if scores is not None and result is not None:
-            recommendations = []
-            for idx in range(len(result['url'])):
-                recommendations.append(
-                    f"Video {idx+1}: {result['url'][idx]}\nPlaycount: {int(result['playCount'][idx])}\n"
-                )
-            recommended_text = "\n\n".join(recommendations)
-            recommended_influencer = f"\nYou can use these influencers for this type of video {str(result['username'][:3])}"
-
-            predictions = self.compute_mean_and_predict(result)
-            generated_prompt = self.generate_prompt(generated_response)
-
-        return recommended_text + recommended_influencer, predictions, generated_prompt
-
-    def format_predictions(self, predictions):
-        """Format predictions for display."""
-        if predictions:
-            return "\n".join([f"{key}: {value}" for key, value in predictions.items()])
-        else:
-            return "No predictions available."
-
-    def launch_interface(self):
-        """Launch the Gradio interface."""
-        interface=gr.Interface(
-            fn=self.generate_story_and_recommendation,
-            inputs=gr.Textbox(label="Enter your generated story.", lines=15),
-            outputs=[
-                gr.Textbox(label="Our Recommendations for you."),
-                gr.Textbox(label="Predicted Metrics (Likes, Comments, Shares)", type="text"),
-                gr.Textbox(label="Recommended Prompt for video generation:"),
-            ],
-            title="Video Story Generation and Recommendation",
-            description="Enter a request for a video storyline, and get a detailed story along with recommended videos and predicted engagement metrics based on the same input."
-        )
-        return interface
+import os
+
+
+import gradio as gr
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer
+from datasets import load_dataset
+import numpy as np
+import torch
+from shared_resources import shared_resources
+from phi.agent import Agent
+from phi.tools.duckduckgo import DuckDuckGo
+from phi.agent import Agent, RunResponse
+from phi.model.huggingface import HuggingFaceChat
+
+
+
+class StoryRecommendationApp:
+
+    def __init__(self):
+        # self.tokenizer = shared_resources.tokenizer
+        # self.model = shared_resources.model
+        self.device = shared_resources.device
+        self.sentence_transformer = shared_resources.sentence_transformer
+        self.data = shared_resources.data
+
+    def search(self, query: str, k: int = 3):
+        """Search for recommended videos based on user input."""
+        embedded_query = self.sentence_transformer.encode(query)
+        scores, retrieved_examples = self.data.get_nearest_examples("embeddings", embedded_query, k=k)
+
+        if isinstance(retrieved_examples, np.ndarray):
+            retrieved_examples = np.nan_to_num(retrieved_examples, nan=0)
+        elif isinstance(retrieved_examples, list):
+            retrieved_examples = [
+                [0 if np.isnan(val) else val for val in example] if isinstance(example, list) else example
+                for example in retrieved_examples
+            ]
+        return scores, retrieved_examples
+
+    def compute_mean_and_predict(self, retrieved_examples):
+        """Compute mean for likesCount, commentCount, and shareCount."""
+        features = ["LikesCount", "commentCount", "shareCount"]
+        predictions = {}
+
+        for feature in features:
+            values = np.array(retrieved_examples[feature])
+            values = np.array([np.nan if v is None else v for v in values])
+            values = np.nan_to_num(values, nan=0.0)
+
+            mean_value = int(np.mean(values))
+            predictions[f"predicted_{feature}"] = mean_value
+
+        return predictions
+
+    def generate_prompt(self, query: str):
+        """Generate a prompt for video generation based on the input story."""
+        input_text = f'''
+        I want to summarize a story in exactly 3 sentences. No more than 3 nor less than 3. 
+        But the sentences have to be good enough to use as a prompt for video generation, because I have to give those 3 sentences to the video generation model.
+        For example: This prompt is about A heartwarming family reunion celebrating love and cherished memories in exactly 3 sentences.
+        -A warm, heartfelt reunion in a cozy living room, where family members embrace each other after a long time apart, soft lighting enhances the intimate atmosphere, and laughter fills the air.
+        -A close-up shot of a grandmother’s hands carefully arranging a family photo album, as the camera pans over old pictures, evoking cherished memories and a deep sense of love.
+        -A final moment around the dinner table, family members sharing a meal together, toasts are made, and the soft glow of candles reflects the joy and connection between generations.
+
+        So, I will provide you that story now. The story is:\n{query}
+
+        Please remember, don't give any background descriptions. Just generate 3 sentences likewise the example above. Don't even give the starting text like: "Here are the 3 sentences that summarize the story:" or any other like this. Just give the answers in 3 sentences directly
+        '''
+        agent = Agent(
+            model=HuggingFaceChat(
+                id="meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8",
+                max_tokens=4096,
+            ),
+            # tools=[DuckDuckGo()],
+            markdown=True
+        )
+
+        # Get the response in a variable
+        run: RunResponse = agent.run(input_text)
+        generated_text=run.content
+
+        sentences = [sentence.strip() for sentence in generated_text.split('.') if sentence]
+        return '. '.join(sentences[-3:]) + ('.' if len(sentences) > 0 else '')
+
+    def generate_story_and_recommendation(self, generated_response: str):
+        """Generate story recommendations and predictions based on the user input."""
+        scores, result = self.search(generated_response, 4)
+
+        recommended_videos_text = ""
+        predictions = {}
+
+        if scores is not None and result is not None:
+            recommendations = []
+            for idx in range(len(result['url'])):
+                recommendations.append(
+                    f"Video {idx+1}: {result['url'][idx]}\nPlaycount: {int(result['playCount'][idx])}\n"
+                )
+            recommended_text = "\n\n".join(recommendations)
+            recommended_influencer = f"\nYou can use these influencers for this type of video {str(result['username'][:3])}"
+
+            predictions = self.compute_mean_and_predict(result)
+            generated_prompt = self.generate_prompt(generated_response)
+
+        return recommended_text + recommended_influencer, predictions, generated_prompt
+
+    def format_predictions(self, predictions):
+        """Format predictions for display."""
+        if predictions:
+            return "\n".join([f"{key}: {value}" for key, value in predictions.items()])
+        else:
+            return "No predictions available."
+
+    def launch_interface(self):
+        """Launch the Gradio interface."""
+        interface=gr.Interface(
+            fn=self.generate_story_and_recommendation,
+            inputs=gr.Textbox(label="Enter your generated story.", lines=15),
+            outputs=[
+                gr.Textbox(label="Our Recommendations for you."),
+                gr.Textbox(label="Predicted Metrics (Likes, Comments, Shares)", type="text"),
+                gr.Textbox(label="Recommended Prompt for video generation:"),
+            ],
+            title="Video Story Generation and Recommendation",
+            description="Enter a request for a video storyline, and get a detailed story along with recommended videos and predicted engagement metrics based on the same input."
+        )
+        return interface