src/summarization.py

# summarization.py
"""
Transcript summarization module with LLM.
Provides a robust function for summarizing long texts using
intelligent chunking and local language models.

Hybrid version: uses LangChain for text splitting and prompts,
but llama_cpp directly for LLM calls (better performance).
"""

import re
import time
from functools import lru_cache
from typing import Iterator
import os
import multiprocessing

from llama_cpp import Llama
# Import from the dedicated langchain packages; the old `langchain.text_splitter`
# / `langchain.prompts` shims were removed in langchain 1.x.
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import PromptTemplate

from .utils import available_gguf_llms, num_vcpus, s2tw_converter

# Detection of available logical cores
detected_cpus = multiprocessing.cpu_count()
is_hf_spaces = os.environ.get('SPACE_ID') is not None
print(f"Detected vCPUs: {detected_cpus}, Effective vCPUs: {num_vcpus}" + (" (HF Spaces limited)" if is_hf_spaces else ""))


# Generation sampling parameters. These are *generation-time* settings and must
# be passed to each create_chat_completion call — llama-cpp-python ignores
# sampling kwargs handed to the Llama constructor, so setting them there had no
# effect on the output. A moderate repeat_penalty plus low temperature keep the
# small local models from looping without over-penalizing legitimate phrasing.
GENERATION_KWARGS = {
    "repeat_penalty": 1.3,
    "temperature": 0.1,
    "top_p": 0.9,
    "top_k": 40,
}


@lru_cache(maxsize=1)
def get_llm(selected_gguf_model: str) -> Llama:
    """Cache and return the LLM model"""
    repo_id, filename = available_gguf_llms[selected_gguf_model]
    return Llama.from_pretrained(
        repo_id=repo_id,
        filename=filename,
        verbose=False,
        n_ctx=4096,
        n_threads=num_vcpus,
    )


def get_language_instruction() -> str:
    """Get universal language instruction for system prompts"""
    return "IMPORTANT: You MUST respond in the EXACT SAME LANGUAGE as the input text. If the input is in French, respond in French. If the input is in Chinese, respond in Chinese. If the input is in Spanish, respond in Spanish. Do not translate to English. Maintain the original language throughout your entire response."


def remove_repetition(text: str, min_length: int = 50) -> str:
    """
    Collapse obvious repetition (the same line or sentence emitted back-to-back)
    in generated text, preserving the original line structure so markdown
    formatting (headings, lists, paragraph breaks) survives.

    This is a light, language-agnostic safeguard only; the primary defense
    against runaway repetition is the repeat_penalty applied at generation time
    (see GENERATION_KWARGS). It never truncates legitimate, non-repetitive
    content — earlier versions hardcoded an English keyword and a blanket
    length cap, which silently discarded valid summaries.
    """
    if not text or len(text) < min_length:
        return text

    def dedupe_sentences(segment: str) -> str:
        # Split into sentences while *keeping* the separators (captured group),
        # so the original spacing is reconstructed exactly. Latin terminators
        # are followed by whitespace; full-width CJM terminators usually are not
        # (zero-width split). This avoids inserting/removing spaces in
        # non-duplicate text.
        tokens = re.split(r'((?<=[.!?])\s+|(?<=[。！？]))', segment)
        result = []
        prev = None
        i = 0
        while i < len(tokens):
            sentence = tokens[i]
            sep = tokens[i + 1] if i + 1 < len(tokens) else ''
            if sentence.strip() and sentence.strip() == prev:
                i += 2  # drop the duplicate sentence and its trailing separator
                continue
            result.append(sentence)
            result.append(sep)
            if sentence.strip():
                prev = sentence.strip()
            i += 2
        return ''.join(result)

    out_lines = []
    for line in text.split('\n'):
        # Collapse a line that exactly repeats the previous non-empty line.
        if out_lines and line.strip() and line.strip() == out_lines[-1].strip():
            continue
        out_lines.append(dedupe_sentences(line) if line.strip() else line)
    return '\n'.join(out_lines).strip()


def get_summarization_instructions() -> str:
    """Get comprehensive summarization instructions to prevent common issues"""
    return """You are an expert transcript summarizer. Create clear, concise summaries that capture key points without ANY repetition.

CRITICAL RULES - NEVER DO THESE:
- NEVER repeat words, phrases, or sentences
- NEVER start with "Here is a summary", "Okay", "Voici un résumé", or similar
- NEVER copy text directly from the input
- NEVER repeat the same ideas multiple times

REQUIRED BEHAVIOR:
- Create NEW, ORIGINAL content that summarizes the main ideas
- Keep summaries concise (aim for 25-35% of original length)
- Focus on 2-3 key points maximum
- Use natural, flowing language
- Be direct and to the point
- Maintain factual accuracy"""


def create_text_splitter(chunk_size: int = 4000, chunk_overlap: int = 200) -> RecursiveCharacterTextSplitter:
    """Create a text splitter with intelligent separators"""
    return RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ". ", " ", ""],
        length_function=len,
    )


def create_chunk_summary_prompt() -> PromptTemplate:
    """Prompt for summarizing an individual chunk"""
    template = """Summarize this part of the transcript while keeping the key points and important information.

Transcript:
{text}

Concise summary:"""
    return PromptTemplate(template=template, input_variables=["text"])


def create_final_summary_prompt() -> PromptTemplate:
    """Prompt for creating the final summary from partial summaries"""
    template = """Here are the summaries of different parts of a transcript.
Create a coherent and synthetic summary of the whole.

{user_prompt}

Partial summaries:
{partial_summaries}

Final summary:"""
    return PromptTemplate(
        template=template,
        input_variables=["user_prompt", "partial_summaries"]
    )


def summarize_chunk(llm: Llama, text: str, prompt_template: PromptTemplate) -> str:
    """Summarize an individual chunk using LangChain for the prompt"""
    try:
        # Use LangChain to format the prompt
        formatted_prompt = prompt_template.format(text=text)

        response = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                {"role": "user", "content": formatted_prompt}
            ],
            stream=False,
            **GENERATION_KWARGS,
        )
        summary = response['choices'][0]['message']['content']
        # Remove repetition from chunk summary
        cleaned_summary = remove_repetition(summary)
        return s2tw_converter.convert(cleaned_summary)
    except Exception as e:
        print(f"Error during chunk summarization: {e}")
        return f"[Summarization error: {str(e)}]"


def summarize_transcript_langchain(transcript: str, selected_gguf_model: str, prompt_input: str) -> Iterator[str]:
    """
    Hybrid LangChain + llama_cpp version of transcript summarization.

    LangChain advantages used:
    - RecursiveCharacterTextSplitter: intelligent chunking with natural separators
    - PromptTemplate: clean template management
    - More readable and maintainable code

    Keeps llama_cpp for LLM calls (better performance).
    """
    if not transcript or not transcript.strip():
        yield "The transcript is empty."
        return

    try:
        # Component initialization
        llm = get_llm(selected_gguf_model)
        text_splitter = create_text_splitter()
        chunk_prompt = create_chunk_summary_prompt()
        final_prompt = create_final_summary_prompt()

        # Token estimation
        transcript_tokens = len(llm.tokenize(transcript.encode('utf-8')))

        # Direct summary if text is short
        if transcript_tokens <= 2000:
            print(f"[summarize_transcript] Direct summary: {transcript_tokens} tokens")

            # Direct completion without streaming to debug
            completion = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                    {"role": "user", "content": f"{prompt_input}\n\n{transcript}"}
                ],
                stream=False,  # Disable streaming to get complete response
                **GENERATION_KWARGS,
            )

            full_response = completion['choices'][0]['message']['content']
            # Remove repetition from the response
            cleaned_response = remove_repetition(full_response)
            yield s2tw_converter.convert(cleaned_response)
            # Short transcript handled directly; do NOT fall through into the
            # chunking/map-reduce path (that would summarize again and overwrite
            # this result on the client).
            return

        # Chunking with LangChain for long texts
        chunks = text_splitter.split_text(transcript)
        print(f"[summarize_transcript] Text divided into {len(chunks)} chunks")

        # Summary of each chunk
        partial_summaries = []
        for i, chunk in enumerate(chunks, 1):
            print(f"Summarizing chunk {i}/{len(chunks)}")
            summary = summarize_chunk(llm, chunk, chunk_prompt)
            partial_summaries.append(summary)

        # Combination and final summary
        combined_summaries = "\n\n".join(partial_summaries)

        # Check combination size
        combined_tokens = len(llm.tokenize(combined_summaries.encode('utf-8')))

        if combined_tokens <= 3500:  # Leave some margin
            print(f"[summarize_transcript] Final summary of {len(partial_summaries)} partial summaries")

            # Use LangChain to format the final prompt
            final_prompt_formatted = final_prompt.format(
                user_prompt=prompt_input,
                partial_summaries=combined_summaries
            )

            # Use non-streaming for final summary to enable repetition removal
            completion = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                    {"role": "user", "content": final_prompt_formatted}
                ],
                stream=False,
                **GENERATION_KWARGS,
            )

            full_response = completion['choices'][0]['message']['content']
            # Remove repetition from the final response
            cleaned_response = remove_repetition(full_response)
            yield s2tw_converter.convert(cleaned_response)
        else:
            print(f"[summarize_transcript] Combination too long ({combined_tokens} tokens), simplified summary")
            # Fallback: direct summary of the combination
            stream = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"{get_summarization_instructions()} {get_language_instruction()}"},
                    {"role": "user", "content": f"{prompt_input}\n\n{combined_summaries}"}
                ],
                stream=True,
                **GENERATION_KWARGS,
            )

            full_response = ""
            for chunk in stream:
                delta = chunk['choices'][0]['delta']
                if 'content' in delta:
                    full_response += delta['content']
                    yield s2tw_converter.convert(full_response)

    except Exception as e:
        print(f"General error during summarization: {e}")
        yield f"[Error during summarization: {str(e)}]"


def create_title_prompt() -> PromptTemplate:
    """Prompt for generating a document title"""
    template = """Generate a single, concise title for this transcript that captures the main topic or theme. Keep it under 10 words.

Transcript:
{text}

Title:"""
    return PromptTemplate(template=template, input_variables=["text"])


def generate_title(transcript: str, selected_gguf_model: str) -> str:
    """
    Generate a title for the transcript using the selected LLM.
    Returns a concise title that captures the main topic.
    """
    if not transcript or not transcript.strip():
        return "Untitled Document"

    try:
        # Get the LLM
        llm = get_llm(selected_gguf_model)
        title_prompt = create_title_prompt()

        # Use first 2000 tokens for title generation to avoid excessive context
        tokens = llm.tokenize(transcript.encode('utf-8'))
        if len(tokens) > 2000:
            # Truncate to first 2000 tokens and decode back to text
            truncated_tokens = tokens[:2000]
            truncated_text = llm.detokenize(truncated_tokens).decode('utf-8')
        else:
            truncated_text = transcript

        # Format the prompt
        formatted_prompt = title_prompt.format(text=truncated_text)

        # Generate title
        response = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": f"You are an expert at creating single, concise titles for documents and transcripts. Always provide exactly one title, nothing else. {get_language_instruction()}"},
                {"role": "user", "content": formatted_prompt}
            ],
            stream=False,
            max_tokens=20,  # Very short for titles
            **GENERATION_KWARGS,
        )

        title = response['choices'][0]['message']['content'].strip()
        # Clean up the title (remove quotes, extra whitespace)
        title = title.strip('"\'').strip()
        # Apply zh-cn to zh-tw conversion
        title = s2tw_converter.convert(title)
        return title if title else "Untitled Document"

    except Exception as e:
        print(f"Error generating title: {e}")
        return "Untitled Document"


def create_speaker_name_detection_prompt() -> PromptTemplate:
    """Prompt for detecting speaker names from their utterances"""
    template = """Analyze the following utterances from a single speaker and suggest a name for this speaker. Look for:

1. Self-introductions or self-references
2. Names mentioned in context
3. Speech patterns, vocabulary, and topics that might indicate identity
4. Professional titles, roles, or relationships mentioned

Utterances from this speaker:
{text}

Based on the content, suggest a name for this speaker. Consider:
- If the speaker introduces themselves, use that name
- If the speaker is addressed by others, use that name
- If the content suggests a clear identity (e.g., "I'm Dr. Smith", "As CEO", "My name is John")
- If no clear name is evident, suggest "Unknown"

Provide your answer in this exact format:
NAME: [suggested name]
CONFIDENCE: [high/medium/low]
REASON: [brief explanation]

If confidence is "low", the name should not be used."""
    return PromptTemplate(template=template, input_variables=["text"])


def detect_speaker_names(utterances: list, selected_gguf_model: str) -> dict:
    """
    Detect speaker names from diarized utterances using LLM analysis.

    Args:
        utterances: List of utterance dicts with 'speaker', 'text', 'start', 'end' keys
        selected_gguf_model: The LLM model to use for analysis

    Returns:
        Dict mapping speaker_id to detected name info:
        {
            speaker_id: {
                'name': str,
                'confidence': str,  # 'high', 'medium', 'low'
                'reason': str
            }
        }
    """
    if not utterances:
        return {}

    # Group utterances by speaker
    speaker_utterances = {}
    for utt in utterances:
        speaker_id = utt.get('speaker')
        if speaker_id is not None:
            text = utt.get('text', '')
            if not text:
                continue
            if speaker_id not in speaker_utterances:
                speaker_utterances[speaker_id] = []
            speaker_utterances[speaker_id].append(text)

    if not speaker_utterances:
        return {}

    try:
        llm = get_llm(selected_gguf_model)
        prompt = create_speaker_name_detection_prompt()

        speaker_names = {}

        for speaker_id, texts in speaker_utterances.items():
            # Combine all utterances for this speaker (limit to reasonable length)
            combined_text = ' '.join(texts)
            if len(combined_text) > 4000:  # Limit context
                combined_text = combined_text[:4000] + '...'

            # Format prompt
            formatted_prompt = prompt.format(text=combined_text)

            # Get LLM response
            response = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": f"You are an expert at analyzing speech patterns and identifying speaker identities from transcripts. Be precise and only suggest names when you have clear evidence. {get_language_instruction()}"},
                    {"role": "user", "content": formatted_prompt}
                ],
                stream=False,
                max_tokens=100,
                **GENERATION_KWARGS,
            )

            result_text = response['choices'][0]['message']['content'].strip()

            # Parse the response
            name = "Unknown"
            confidence = "low"
            reason = "No clear identification found"

            lines = result_text.split('\n')
            for line in lines:
                if line.startswith('NAME:'):
                    name = line.replace('NAME:', '').strip()
                elif line.startswith('CONFIDENCE:'):
                    confidence = line.replace('CONFIDENCE:', '').strip().lower()
                elif line.startswith('REASON:'):
                    reason = line.replace('REASON:', '').strip()

            # Only include high confidence detections
            if confidence == 'high' and name != "Unknown":
                speaker_names[speaker_id] = {
                    'name': name,
                    'confidence': confidence,
                    'reason': reason
                }

        return speaker_names

    except Exception as e:
        print(f"Error detecting speaker names: {e}")
        return {}


# Alias pour maintenir la compatibilité
summarize_transcript = summarize_transcript_langchain