Update main.py

main.py

@@ -1,15 +1,40 @@
+import gradio as gr
 import whisper
-
-def transcribe_audio(audio_path):
-    model = whisper.load_model("base")
-    result = model.transcribe(audio_path)
-    return result["text"]
 from pyannote.audio import Pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from bertopic import BERTopic
+from sklearn.feature_extraction.text import CountVectorizer
+import pandas as pd
+import torch
+
+# Load Whisper model for transcription
+whisper_model = whisper.load_model("large")
+
+# Load translation pipeline
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ar-en")
+
+# Load summarization pipeline
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+
+# Load LLaMA model and tokenizer for chat-based interaction
+llama_model_name = "meta-llama/Llama-2-7b-chat"
+tokenizer = AutoTokenizer.from_pretrained(llama_model_name)
+model = AutoModelForCausalLM.from_pretrained(llama_model_name)
+
+# Global variables to store processed data
+aligned_transcription = []
+translated_text = ""
+topics = []
+summary = ""
 
-def perform_speaker_diarization(audio_path):
-    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization@2.1", use_auth_token="YOUR_HUGGINGFACE_TOKEN")
+def perform_speaker_diarization(audio_path, hf_token="YOUR_HUGGINGFACE_TOKEN"):
+    # Load the speaker diarization pipeline
+    pipeline = Pipeline.from_pretrained("pyannote/[email protected]", use_auth_token=hf_token)
+    
+    # Apply diarization
     diarization = pipeline(audio_path)
     
+    # Extract speaker segments
     speaker_segments = []
     for turn, _, speaker in diarization.itertracks(yield_label=True):
         speaker_segments.append({
@@ -19,145 +44,108 @@ def perform_speaker_diarization(audio_path):
         })
     return speaker_segments
 
-from textblob import TextBlob
-from sklearn.feature_extraction.text import CountVectorizer
-from sklearn.decomposition import LatentDirichletAllocation
-from collections import Counter
-import nltk
-from nltk.corpus import stopwords
-import spacy
-
-nltk.download('stopwords')
-nltk.download('punkt')
-
-# Load spaCy model for NER
-nlp = spacy.load("en_core_web_sm")
-
-def analyze_sentiment(text):
-    blob = TextBlob(text)
-    return blob.sentiment.polarity, blob.sentiment.subjectivity
-
-def extract_keywords(text, top_n=5):
-    stop_words = set(stopwords.words("english"))
-    words = nltk.word_tokenize(text.lower())
-    filtered_words = [word for word in words if word.isalnum() and word not in stop_words]
-    word_counts = Counter(filtered_words)
-    return word_counts.most_common(top_n)
+def transcribe_with_speaker_diarization(audio_path, hf_token="YOUR_HUGGINGFACE_TOKEN"):
+    # Step 1: Perform speaker diarization
+    speaker_segments = perform_speaker_diarization(audio_path, hf_token)
+    
+    # Step 2: Transcribe audio
+    transcription = whisper_model.transcribe(audio_path)
+    
+    # Step 3: Align transcription with speaker segments
+    aligned_transcription = []
+    for segment in transcription["segments"]:
+        start_time = segment["start"]
+        end_time = segment["end"]
+        text = segment["text"]
+        
+        # Find the corresponding speaker
+        speaker = "Unknown"
+        for spk_segment in speaker_segments:
+            if spk_segment["start"] <= start_time <= spk_segment["end"]:
+                speaker = spk_segment["speaker"]
+                break
+        
+        aligned_transcription.append({
+            "speaker": speaker,
+            "start": start_time,
+            "end": end_time,
+            "text": text
+        })
+    
+    return aligned_transcription
 
-def perform_topic_modeling(text, num_topics=5, num_words=10):
-    vectorizer = CountVectorizer(stop_words="english", max_features=1000)
-    X = vectorizer.fit_transform([text])
-    lda = LatentDirichletAllocation(n_components=num_topics, random_state=42)
-    lda.fit(X)
+def translate_text(text, src_lang="ar", tgt_lang="en"):
+    translated = translator(text, max_length=400)
+    return translated[0]["translation_text"]
 
-    topics = []
-    for idx, topic in enumerate(lda.components_):
-        top_words = [vectorizer.get_feature_names_out()[i] for i in topic.argsort()[:-num_words - 1:-1]]
-        topics.append(f"Topic {idx + 1}: {' '.join(top_words)}")
-    return topics
+def perform_topic_modeling(texts):
+    vectorizer = CountVectorizer(stop_words="english")
+    topic_model = BERTopic(vectorizer_model=vectorizer, calculate_probabilities=True)
+    topics, probs = topic_model.fit_transform(texts)
+    return topic_model.get_topic_info(), topic_model.visualize_topics()
 
-def extract_entities(text):
-    doc = nlp(text)
-    entities = [(ent.text, ent.label_) for ent in doc.ents]
-    return entities
+def summarize_text(text, max_length=150, min_length=30):
+    summary = summarizer(text, max_length=max_length, min_length=min_length, do_sample=False)
+    return summary[0]["summary_text"]
 
-def parse_query(query):
-    doc = nlp(query)
-    keywords = [token.text.lower() for token in doc if token.is_alpha and not token.is_stop]
-    intent = None
-    
-    if any(word in ["how many", "count"] for word in keywords):
-        intent = "count"
-    elif any(word in ["list", "show me"] for word in keywords):
-        intent = "list"
-    elif any(word in ["sentiment", "polarity", "subjectivity"] for word in keywords):
-        intent = "sentiment"
-    elif any(word in ["theme", "topic", "main"] for word in keywords):
-        intent = "topic"
-    elif any(word in ["keyword", "common"] for word in keywords):
-        intent = "keyword"
-    elif any(word in ["entity", "name", "person", "organization"] for word in keywords):
-        intent = "ner"
-    return intent, keywords
+def generate_response(prompt, max_tokens=150):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=max_tokens)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
-def answer_question(query, qa_df):
-    intent, keywords = parse_query(query)
-    
-    if intent == "count":
-        filtered = qa_df[qa_df["Transcript"].str.contains("|".join(keywords), case=False)]
-        return f"{len(filtered)} responses contain the keywords: {', '.join(keywords)}."
-    
-    elif intent == "list":
-        filtered = qa_df[qa_df["Transcript"].str.contains("|".join(keywords), case=False)]["Transcript"].tolist()
-        return "\n".join(filtered) if filtered else "No matching responses found."
+def process_audio(audio_path, language="auto", hf_token="YOUR_HUGGINGFACE_TOKEN"):
+    global aligned_transcription, translated_text, topics, summary
     
-    elif intent == "sentiment":
-        avg_polarity = qa_df["Sentiment_Polarity"].mean()
-        avg_subjectivity = qa_df["Sentiment_Subjectivity"].mean()
-        return f"Average Polarity: {avg_polarity:.2f}, Average Subjectivity: {avg_subjectivity:.2f}"
-    
-    elif intent == "topic":
-        all_text = " ".join(qa_df["Transcript"])
-        topics = perform_topic_modeling(all_text)
-        return "\n".join(topics)
-    
-    elif intent == "keyword":
-        all_text = " ".join(qa_df["Transcript"])
-        keywords = extract_keywords(all_text)
-        return ", ".join([word for word, count in keywords])
-    
-    elif intent == "ner":
-        all_text = " ".join(qa_df["Transcript"])
-        entities = extract_entities(all_text)
-        return "\n".join([f"{entity} ({label})" for entity, label in entities])
+    # Step 1: Transcribe audio with speaker diarization
+    aligned_transcription = transcribe_with_speaker_diarization(audio_path, hf_token)
     
+    # Step 2: Translate text if needed
+    full_text = " ".join([seg["text"] for seg in aligned_transcription])
+    if language != "en":
+        translated_text = translate_text(full_text, src_lang="ar", tgt_lang="en")
     else:
-        return "I'm not sure how to answer that. Try asking about counts, lists, sentiment, topics, keywords, or entities."
-
-import gradio as gr
-
-# Global variables to store processed data
-qa_df = None
-
-def process_audio(audio_path):
-    global qa_df
+        translated_text = full_text
     
-    # Step 1: Transcribe audio
-    transcription = transcribe_audio(audio_path)
+    # Step 3: Perform topic modeling
+    topics, _ = perform_topic_modeling([translated_text])
     
-    # Step 2: Perform speaker diarization
-    speaker_segments = perform_speaker_diarization(audio_path)
+    # Step 4: Summarize text
+    summary = summarize_text(translated_text)
     
-    # Step 3: Analyze text
-    sentiment_polarity, sentiment_subjectivity = analyze_sentiment(transcription)
-    topics = perform_topic_modeling(transcription)
-    keywords = extract_keywords(transcription)
-    entities = extract_entities(transcription)
+    return "Audio processed successfully!"
+
+def answer_question(query):
+    global aligned_transcription, translated_text, topics, summary
     
-    # Create a DataFrame
-    qa_df = pd.DataFrame({
-        "Speaker": [seg["speaker"] for seg in speaker_segments],
-        "Transcript": [transcription],
-        "Sentiment_Polarity": [sentiment_polarity],
-        "Sentiment_Subjectivity": [sentiment_subjectivity],
-        "Topics": [topics],
-        "Keywords": [keywords],
-        "Entities": [entities]
-    })
+    # Combine context for the LLM
+    context = f"""
+    Transcription: {translated_text}
+    Topics: {topics.to_string(index=False)}
+    Summary: {summary}
+    """
     
-    return "Audio processed successfully!"
+    # Generate response using LLM
+    response = generate_response(f"{context}\nQuestion: {query}")
+    return response
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Advanced Audio Analysis App")
+    gr.Markdown("# Advanced Audio Analysis App with Speaker Diarization")
     audio_input = gr.Audio(label="Upload Audio File")
+    language_input = gr.Dropdown(choices=["auto", "en", "ar"], label="Language", value="auto")
+    hf_token_input = gr.Textbox(label="Hugging Face Token (for pyannote.audio)", type="password")
     process_button = gr.Button("Process Audio")
     status_output = gr.Textbox(label="Status")
     
     question_input = gr.Textbox(label="Ask a Question")
     answer_output = gr.Textbox(label="Answer")
     
-    process_button.click(process_audio, inputs=audio_input, outputs=status_output)
-    question_input.submit(answer_question, inputs=[question_input], outputs=answer_output)
+    process_button.click(
+        process_audio,
+        inputs=[audio_input, language_input, hf_token_input],
+        outputs=status_output
+    )
+    question_input.submit(answer_question, inputs=question_input, outputs=answer_output)
 
 demo.launch()