pipline.py

import argparse
import os
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
import faiss
import gc
import openai
import random
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# ==============================
# Part 1: Generate RAG Prompt and Save to File
# ==============================
def set_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        

def generate_rag_prompt(user_prompt, 
                        output_dir, 
                        total_videos=1974, 
                        selected_videos_num=50, 
                        num_tags=1, 
                        ratio=0.1,  # ratio represents the proportion of negative samples
                        seed=42):
    # Data Loading
    set_random_seed(seed)

    # 2. Load various data files
    views_data = pd.read_csv('Microlens/MicroLens-100k_likes_and_views.txt',
                             sep='\t', header=None, names=['video_id','likes','views'])
    # MicroLens-100k_title_en.csv contains cover image descriptions (cover_desc)
    cover_data = pd.read_csv('Microlens/MicroLens-100k_title_en.csv',
                             sep=',', header=None, names=['video_id','cover_desc'],
                             on_bad_lines='skip')
    # Microlens100K_captions_en.csv contains video captions; use as title_en and caption_en
    desc_data = pd.read_csv('Microlens/Microlens100K_captions_en.csv',
                            sep='\t', header=None, names=['video_id','caption_en'],
                            on_bad_lines='skip')
    tags_data = pd.read_csv('Microlens/tags_to_summary.csv',
                            sep=',', header=None, names=['video_id','partition'],
                            on_bad_lines='skip')

    # 3. Load comment data and count comments for each video_id
    comments_data = pd.read_csv('Microlens/MicroLens-100k_comment_en.txt',
                                sep='\t', header=None, names=['user_id','video_id','comment_text'],
                                on_bad_lines='skip')
    comments_data = comments_data[['video_id','comment_text']]

    # Group by video_id and count comments
    comment_count_df = (
        comments_data
        .groupby('video_id')['comment_text']
        .count()
        .reset_index(name='comment_count')
    )

    # Merge all data (use cover_desc as title_en since no separate title file is available)
    merged = (
        views_data
        .merge(cover_data, on='video_id', how='left')
        .merge(desc_data, on='video_id', how='left')
        .merge(tags_data, on='video_id', how='left')
        .merge(comment_count_df, on='video_id', how='left')
    )
    # Use cover_desc as title_en for compatibility with downstream prompts
    merged['title_en'] = merged['cover_desc']

    # Drop rows with missing values in key fields
    merged.dropna(subset=['title_en', 'cover_desc', 'caption_en', 'partition', 'comment_count'], inplace=True)

    # 4. Use SentenceTransformer to create embeddings for 'partition' and user prompts
    model = SentenceTransformer('sentence-transformers/all-MiniLM-L12-v2')
    unique_partitions = merged['partition'].unique().tolist()
    partition_embeddings = model.encode(unique_partitions)
    prompt_embedding = model.encode([user_prompt])[0]

    # Calculate similarities and get top partitions
    similarities = [
        np.dot(prompt_embedding, pe) 
        / (np.linalg.norm(prompt_embedding) * np.linalg.norm(pe)) 
        for pe in partition_embeddings
    ]
    top_k_indices = np.argsort(similarities)[::-1][:num_tags]
    top_k_partitions = [unique_partitions[i] for i in top_k_indices]

    # Filter data by top partitions
    filtered_data = merged[merged['partition'].isin(top_k_partitions)]

    # Sort by comment count
    filtered_data = filtered_data.sort_values('comment_count', ascending=False)

    # Split dataset based on ratio (proportion of negative samples)
    n_negative = int(len(filtered_data) * ratio)
    n_positive = len(filtered_data) - n_negative

    # Select popular videos (high comment count)
    positive_videos = filtered_data.head(n_positive)

    # Select unpopular videos from remaining data
    remaining_data = filtered_data.iloc[n_positive:]
    negative_videos = remaining_data.tail(n_negative)

    # Remove duplicates to prevent overlap
    positive_videos.drop_duplicates(subset=['video_id'], inplace=True)
    negative_videos.drop_duplicates(subset=['video_id'], inplace=True)

    # Merge positive and negative samples for retrieval
    combined_videos = pd.concat([positive_videos, negative_videos])
    combined_videos.drop_duplicates(subset=['video_id'], keep='first', inplace=True)

    # Create combined text for embedding
    combined_texts = (
        combined_videos['title_en'] + " " 
        + combined_videos['cover_desc'] + " " 
        + combined_videos['caption_en']
    ).tolist()
    combined_embeddings = model.encode(combined_texts)

    # Perform similarity search
    index = faiss.IndexFlatL2(combined_embeddings.shape[1])
    index.add(np.array(combined_embeddings).astype('float32'))
    query_embedding = model.encode([user_prompt]).astype('float32')
    _, I = index.search(query_embedding, len(combined_videos))
    retrieved_videos = combined_videos.iloc[I[0]]

    # Calculate final sample sizes
    final_n_negative = int(selected_videos_num * ratio)
    final_n_positive = selected_videos_num - final_n_negative

    # Find positive and negative samples from retrieval results
    positive_ids = set(positive_videos['video_id'].tolist())
    negative_ids = set(negative_videos['video_id'].tolist())

    retrieved_positive = retrieved_videos[retrieved_videos['video_id'].isin(positive_ids)]
    retrieved_negative = retrieved_videos[retrieved_videos['video_id'].isin(negative_ids)]

    # Select final samples
    final_selected_positive = retrieved_positive.head(final_n_positive)
    final_selected_negative = retrieved_negative.head(final_n_negative)

    # Maintain retrieval order
    final_selected_videos = pd.concat([final_selected_positive, final_selected_negative])
    final_selected_videos = final_selected_videos.loc[
        retrieved_videos.index.intersection(final_selected_videos.index)
    ]
    # Build output text
    rag_positive_context = "\n".join([
        f"Reference Video {i+1} (Positive Sample - Popular):\n"
        f"Title: {row['title_en']}\n"
        f"Cover Desc: {row['cover_desc']}\n"
        f"Desc: {row['caption_en']}\n"
        f"Comment Count: {row['comment_count']}\n"
        for i, (idx, row) in enumerate(final_selected_positive.iterrows())
    ])

    rag_negative_context = "\n".join([
        f"Reference Video {i+1} (Negative Sample - Unpopular):\n"
        f"Title: {row['title_en']}\n"
        f"Cover Desc: {row['cover_desc']}\n"
        f"Desc: {row['caption_en']}\n"
        f"Comment Count: {row['comment_count']}\n"
        for i, (idx, row) in enumerate(final_selected_negative.iterrows())
    ])

    rag_context = rag_positive_context + "\n" + rag_negative_context
    
    cot_prompt = f"""
    Now that you're a talented video creator with a wealth of ideas, you need to think from the user's perspective and after that generate the most popular video title, an AI-generated cover prompt, and a 3-second AI-generated video prompt.
    
    Below is the user query:
    
    {user_prompt}
    
    Below is the reasoning chain (Chain of Thought) that you should follow step-by-step.
    
    Reasoning Chain:
    1. Analyze both popular and unpopular videos as references using the provided context:
       {rag_context}   
    2. Based on the analyzed videos, brainstorm unique and creative ideas for a new video topic. Ensure that the idea is original and does not replicate existing content, as direct copying is strictly forbidden.
    3. Write or conceptualize a logical and original script or content based on a Step2 topic
    4. Double-check:
       1. Whether the theme and content are accurately conveyed
       2. Whether the theme and content are strongly related and complete in fulfilling the user's needs
    5. Start generating based on confirmed topics and content
    6. Re-evaluate:
       1. Whether the generated prompt is logically correct.
       2. Check if the final suggestions match popular trends
    7. If it doesn't meet expectations, refine and finalize the output.  
    
    Explicitly generate a chain of thought during the reasoning process for each candidate. The chain of thought should detail the steps, considerations, and rationale behind the candidate generation, ensuring transparency and clarity in the decision-making process.
    
    Final Answer Requirements:
    - A single line for the final generated Title (MAX_length = 50).
    - A single paragraph for the Cover Prompt.
    - A single paragraph for the Video Prompt (3-second).
    
    
    Now, based on the above reasoning, generate the response in JSON format， here is an example:
    {{
      "title": "Unveiling the Legacy of Ancient Rome: Rise, Glory, and Downfall.",
      "cover_prompt": "Generate an image of a Roman Emperor standing proudly in front of the Colosseum, with a subtle sunset backdrop, highlighting the contrast between the ancient structure.",
      "video_prompt": "Open with a 3-second aerial shot of the Roman Forum, showcasing the sprawling ancient ruins against a clear blue sky, before zooming in on a singular, imposing structure like the Colosseum."
    }}
    Please provide your answer following this exact JSON template for the response.
    """
    
    # Save prompt to file
    output_file = os.path.join(output_dir, f"{user_prompt}.txt")
    with open(output_file, "w", encoding="utf-8") as f:
        f.write(cot_prompt)
    
    print(f"RAG prompt saved to {output_file}")



# ==============================
# Part 2: Load RAG Prompts and Perform Inference
# ==============================
def inference_from_prompts(input_dir, output_dir,seed=42):
    set_random_seed(seed)
    llama_model_name = "meta-llama/Llama-3.3-70B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(llama_model_name)
    quantization_config = BitsAndBytesConfig(load_in_4bit=True)
    model_llama = AutoModelForCausalLM.from_pretrained(
        llama_model_name, device_map="auto", torch_dtype=torch.bfloat16, quantization_config=quantization_config
    )

    llama_pipeline = pipeline(
        "text-generation",
        model=model_llama,
        tokenizer=tokenizer,
        max_new_tokens=5000,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        do_sample=True
    )
    
    for file_name in tqdm(os.listdir(input_dir), desc="inferencing"):
        output_path = os.path.join(output_dir, file_name)
        if os.path.exists(output_path):
            continue
            
        if file_name.endswith(".txt"):
            # 读取输入文件内容
            input_path = os.path.join(input_dir, file_name)
            with open(input_path, "r", encoding="utf-8") as f:
                cot_prompt = f.read()
    
            # 定义 system 和 user 的内容
            messages = [
                {"role": "system", "content": "Now that you're a talented video creator with a wealth of ideas, you need to think from the user's perspective and after that generate the most popular video title, an AI-generated cover prompt, and a 3-second AI-generated video prompt."},
                {"role": "user", "content": cot_prompt},
            ]
            
    
            # 调用模型推理
            response = llama_pipeline(messages, num_return_sequences=1)
            full_output = response[0]['generated_text']
    
            # 保存输出到指定路径
            output_path = os.path.join(output_dir, file_name)
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(full_output[2]['content'])
    
            print(f"Inference result saved to {output_path}")

def inference_from_prompts_qwen(input_dir, output_dir, seed=42):
    # Set random seed for reproducibility
    set_random_seed(seed)
    
    # Qwen model name
    qwen_model_name = "Qwen/Qwen2.5-72B-Instruct"
    
    # Load tokenizer with trust_remote_code=True
    tokenizer = AutoTokenizer.from_pretrained(qwen_model_name, trust_remote_code=True)
    
    # Configure 4-bit quantization
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
    )
    
    # Load Qwen model with auto device mapping and quantization config
    model_qwen = AutoModelForCausalLM.from_pretrained(
        qwen_model_name,
        device_map="auto",
        quantization_config=bnb_config,
        trust_remote_code=True
    )
    
    # Create text generation pipeline with consistent parameters
    qwen_pipeline = pipeline(
        "text-generation",
        model=model_qwen,
        tokenizer=tokenizer,
        max_new_tokens=5000,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        do_sample=True
    )
    
    # Process all txt files in input directory
    for file_name in tqdm(os.listdir(input_dir), desc="inferencing"):
        if file_name.endswith(".txt"):
            # Read input file content
            input_path = os.path.join(input_dir, file_name)
            with open(input_path, "r", encoding="utf-8") as f:
                cot_prompt = f.read()
            
            # Define system and user messages
            messages = [
                {
                    "role": "system", 
                    "content": (
                        "Now that you're a talented video creator with a wealth of ideas, "
                        "you need to think from the user's perspective and after that generate "
                        "the most popular video title, an AI-generated cover prompt, and a 3-second "
                        "AI-generated video prompt."
                    )
                },
                {"role": "user", "content": cot_prompt},
            ]
            
            # Call model for inference
            response = qwen_pipeline(messages, num_return_sequences=1)
            full_output = response[0]['generated_text']
            
            # Save output to specified path
            output_path = os.path.join(output_dir, file_name)
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(full_output[2]['content'])
            
            print(f"Inference result saved to {output_path}")

def inference_from_prompts_mistral(input_dir, output_dir, seed=42):
    # Set random seed for reproducibility
    set_random_seed(seed)
    
    # Mistral model name
    mistral_model_name = "mistralai/Mistral-Large-Instruct-2411"
    
    # Load tokenizer with trust_remote_code=True
    tokenizer = AutoTokenizer.from_pretrained(mistral_model_name, trust_remote_code=True)
    
    # Configure 4-bit quantization
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
    )
    
    # Load Mistral model with auto device mapping and quantization config
    model_mistral = AutoModelForCausalLM.from_pretrained(
        mistral_model_name,
        device_map="auto",
        quantization_config=bnb_config,
        trust_remote_code=True
    )
    
    # Create text generation pipeline with consistent parameters
    mistral_pipeline = pipeline(
        "text-generation",
        model=model_mistral,
        tokenizer=tokenizer,
        max_new_tokens=5000,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        do_sample=True
    )
    
    # Process all txt files in input directory
    for file_name in tqdm(os.listdir(input_dir), desc="inferencing"):
        output_path = os.path.join(output_dir, file_name)
        if os.path.exists(output_path):
            continue
            
        if file_name.endswith(".txt"):
            # Read input file content
            input_path = os.path.join(input_dir, file_name)
            with open(input_path, "r", encoding="utf-8") as f:
                cot_prompt = f.read()
            
            # Define system and user messages
            messages = [
                {
                    "role": "system", 
                    "content": (
                        "Now that you're a talented video creator with a wealth of ideas, "
                        "you need to think from the user's perspective and after that generate "
                        "the most popular video title, an AI-generated cover prompt, and a 3-second "
                        "AI-generated video prompt."
                    )
                },
                {"role": "user", "content": cot_prompt},
            ]
            
            # Call model for inference
            response = mistral_pipeline(messages, num_return_sequences=1)
            full_output = response[0]['generated_text']
            
            # Save output to specified path
            output_path = os.path.join(output_dir, file_name)
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(full_output[2]['content'])
            
            print(f"Inference result saved to {output_path}")

import os
import time
from tqdm import tqdm

def inference_from_prompts_api(input_dir, output_dir, api_key, model="deepseek-reasoner", base_url="https://api.deepseek.com"):
    # Recommended to use 'import openai' directly. Adjust based on your actual library name
    from openai import OpenAI 
    print(model)
    
    # Initialize DeepSeek client
    client = OpenAI(api_key=api_key, base_url=base_url)

    # Iterate through files in input directory
    for file_name in tqdm(os.listdir(input_dir), desc="Inferencing with DeepSeek"):
        output_path = os.path.join(output_dir, file_name)
        if file_name.endswith(".txt") and not os.path.exists(output_path):
            # Read input file content
            input_path = os.path.join(input_dir, file_name)
            with open(input_path, "r", encoding="utf-8") as f:
                cot_prompt = f.read()

            # Define system and user content
            messages = [
                {
                    "role": "system", 
                    "content": (
                        "Now that you're a talented video creator with a wealth of ideas, "
                        "you need to think from the user's perspective and after that generate "
                        "the most popular video title, an AI-generated cover prompt, and a 3-second "
                        "AI-generated video prompt."
                    )
                },
                {"role": "user", "content": cot_prompt},
            ]
            
            # Retry mechanism
            while True:
                try:
                    # Call DeepSeek API for inference
                    response = client.chat.completions.create(
                        model=model,
                        messages=messages
                    )
                    # Break loop if successful
                    break
                except Exception as e:
                    # Print error message and wait 60 seconds before retrying
                    print(f"Error occurred: {e}")
                    print("Waiting 60 seconds before retrying...")
                    time.sleep(120)

            full_output = response.choices[0].message.content

            output_path = os.path.join(output_dir, file_name)
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(full_output)

            print(f"Inference result saved to {output_path}")



def inference_from_prompts_gpt(input_dir, output_dir, api_key, model="gpt-4o"):
    """
    Use OpenAI GPT API to process input prompts and generate outputs.

    Parameters:
        input_dir (str): Directory containing input text files with prompts.
        output_dir (str): Directory to save the generated outputs.
        api_key (str): OpenAI API key for authentication.
        model (str): The GPT model to use (default: "gpt-4").
    """
    # Set the OpenAI API key
    openai.api_key = api_key

    # Ensure the output directory exists
    os.makedirs(output_dir, exist_ok=True)

    # Iterate through files in the input directory
    for file_name in tqdm(os.listdir(input_dir), desc="Inferencing with GPT"):
    
        if file_name.endswith(".txt"):
            # Read input file content
            input_path = os.path.join(input_dir, file_name)
            with open(input_path, "r", encoding="utf-8") as f:
                cot_prompt = f.read()

            # Define the system and user prompts
            messages = [
                {"role": "system", "content": "You are a talented video creator with a wealth of ideas. Think from the user's perspective and generate the most popular video title, an AI-generated cover prompt, and a 3-second AI-generated video prompt."},
                {"role": "user", "content": cot_prompt},
            ]

            # Call the GPT API for inference
            try:
                response = openai.ChatCompletion.create(
                    model=model,
                    messages=messages,
                    stream=False
                )
                # Extract the generated content
                full_output = response['choices'][0]['message']['content']

                # Save the output to the specified path
                output_path = os.path.join(output_dir, file_name)
                with open(output_path, "w", encoding="utf-8") as f:
                    f.write(full_output)

                print(f"Inference result saved to {output_path}")

            except Exception as e:
                print(f"Error processing {file_name}: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--USER_PROMPT", type=str, help="The user prompt to process.")
    parser.add_argument("--MODE", type=str, choices=["generate", "infer","infer_api","infer_gpt","infer_qwen","infer_mistral"], required=True, help="Mode: generate or infer.")
    parser.add_argument("--INPUT_DIR", type=str, help="Input directory for inference.")
    parser.add_argument("--OUTPUT_DIR", type=str, required=True, help="Output directory.")
    parser.add_argument("--MODEL", type=str, required=True, help="Model.")
    parser.add_argument("--VID_NUM", type=int, default=10,help="number of selected videos")
    parser.add_argument("--TAGS_NUM", type=int, default=1,help="number of selected videos")
    args = parser.parse_args()

    if args.MODE == "generate":
        os.makedirs(args.OUTPUT_DIR, exist_ok=True)
        generate_rag_prompt(args.USER_PROMPT, args.OUTPUT_DIR, selected_videos_num=args.VID_NUM,num_tags=args.TAGS_NUM)
    elif args.MODE == "infer":
        os.makedirs(args.OUTPUT_DIR, exist_ok=True)
        inference_from_prompts(args.INPUT_DIR, args.OUTPUT_DIR)
    elif args.MODE == "infer_qwen":
        os.makedirs(args.OUTPUT_DIR, exist_ok=True)
        inference_from_prompts_qwen(args.INPUT_DIR, args.OUTPUT_DIR)
    elif args.MODE == "infer_mistral":
        os.makedirs(args.OUTPUT_DIR, exist_ok=True)
        inference_from_prompts_mistral(args.INPUT_DIR, args.OUTPUT_DIR)
    elif args.MODE == "infer_api":
        os.makedirs(args.OUTPUT_DIR, exist_ok=True)
        inference_from_prompts_api(args.INPUT_DIR, args.OUTPUT_DIR,"API_KEY",model=args.MODEL)
    elif args.MODE == "infer_gpt":
        api_key = "API_KEY"
        os.makedirs(args.OUTPUT_DIR, exist_ok=True)
        inference_from_prompts_gpt(args.INPUT_DIR, args.OUTPUT_DIR,api_key,model="gpt-4o")
    import torch
    
    gc.collect()
    torch.cuda.empty_cache()