import nltk
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.cluster import KMeans
from transformers import pipeline
from nltk.tokenize import sent_tokenize
nltk.download('punkt')

# Load the sentence transformer model for embeddings
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
summarizer = pipeline('summarization')

def embed_questions(questions):
    embeddings = embedding_model.encode(questions, show_progress_bar=True)
    return embeddings

def cluster_questions(embeddings, num_clusters):
    kmeans = KMeans(n_clusters=num_clusters, random_state=42)
    kmeans.fit(embeddings)
    return kmeans.labels_

def summarize_cluster(questions_in_cluster):

    combined_text = ' '.join(questions_in_cluster) # Combine all questions into a single text
    sentences = sent_tokenize(combined_text) # Tokenize into sentences to ensure the summarizer has coherent input

    # Due to model input size limits, process in chunks if necessary
    max_chunk_size = 500  # Adjust based on the summarizer's token limit
    chunks = []
    current_chunk = ''

    for sentence in sentences:
        if len(current_chunk) + len(sentence) <= max_chunk_size:
            current_chunk += ' ' + sentence
        else:
            chunks.append(current_chunk)
            current_chunk = sentence

    if current_chunk:
        chunks.append(current_chunk)

    # Generate summaries for each chunk and combine them
    summaries = []
    for chunk in chunks:
        summary = summarizer(chunk, max_length=60, min_length=30, do_sample=False)
        summaries.append(summary[0]['summary_text'])

    # Combine summaries into a final summary
    final_summary = ' '.join(summaries)
    return final_summary

def process_questions(questions, num_clusters=5):
    # Step 1: Embed the questions
    embeddings = embed_questions(questions)

    # Step 2: Cluster the questions
    labels = cluster_questions(embeddings, num_clusters)

    # Organize questions by their cluster labels
    clustered_questions = {}
    for idx, label in enumerate(labels):
        if label not in clustered_questions:
            clustered_questions[label] = []
        clustered_questions[label].append(questions[idx])

    # Step 3: Summarize each cluster
    cluster_summaries = {}
    for label, qs in clustered_questions.items():
        summary = summarize_cluster(qs)
        cluster_summaries[label] = {
            'summary': summary,
            'questions': qs
        }

    return cluster_summaries


if __name__ == "__main__":
    # Sample list of questions
    # You may modify this to read all the text questions from excel and pass them into a list
    questions = [
        "How can I reset my account password?",
        "What are the pricing plans available?",
        "How do I upgrade my subscription?",
        "I forgot my password, how to retrieve it?",
        "Can I get a discount on bulk purchases?",
        "What is the refund policy?",
        "How to contact customer support?",
        "Is there a free trial available?",
        "How do I change my billing information?",
        "What payment methods are accepted?"
    ]

    num_clusters = 3  # Define the number of clusters based on data size and diversity
    results = process_questions(questions, num_clusters)

    # Display the summaries and associated questions
    for label, content in results.items():
        print(f"\nCluster {label + 1} Summary:")
        print(content['summary'])
        print("\nAssociated Questions:")
        for q in content['questions']:
            print(f"- {q}")

Write a simple embeddings