Create app.py

app.py

@@ -0,0 +1,132 @@
+import torch
+from transformers import pipeline
+import librosa
+from datetime import datetime
+from deep_translator import GoogleTranslator
+from typing import Dict, Union
+from gliner import GLiNER
+import gradio as gr
+
+# Load transcription models
+whisper_pipeline_agri = pipeline("automatic-speech-recognition", model="maliahson/whisper-agri")
+device = 0 if torch.cuda.is_available() else "cpu"
+
+# Initialize GLiNER for information extraction
+gliner_model = GLiNER.from_pretrained("xomad/gliner-model-merge-large-v1.0").to("cpu")
+
+def merge_entities(entities):
+    if not entities:
+        return []
+    merged = []
+    current = entities[0]
+    for next_entity in entities[1:]:
+        if next_entity['entity'] == current['entity'] and (next_entity['start'] == current['end'] + 1 or next_entity['start'] == current['end']):
+            current['word'] += ' ' + next_entity['word']
+            current['end'] = next_entity['end']
+        else:
+            merged.append(current)
+            current = next_entity
+    merged.append(current)
+    return merged
+
+def transcribe_audio(audio_path):
+    """
+    Transcribe a local audio file using the Whisper pipeline, log timing, and save transcription to a file.
+    """
+    try:
+        # Log start time
+        start_time = datetime.now()
+
+        # Ensure audio is mono and resampled to 16kHz
+        audio, sr = librosa.load(audio_path, sr=16000, mono=True)
+
+        # Perform transcription
+        transcription = whisper_pipeline_agri(audio, batch_size=8)["text"]
+
+        # Log end time
+        end_time = datetime.now()
+
+        return transcription
+
+    except Exception as e:
+        return f"Error processing audio: {e}"
+
+def translate_text_to_english(text):
+    """
+    Translate text into English using GoogleTranslator.
+    """
+    try:
+        # Perform translation
+        translated_text = GoogleTranslator(source='auto', target='en').translate(text)
+        return translated_text
+    except Exception as e:
+        return f"Error during translation: {e}"
+
+def extract_information(prompt: str, text: str, threshold: float, nested_ner: bool) -> Dict[str, Union[str, int, float]]:
+    """
+    Extract entities from the English text using GLiNER model.
+    """
+    try:
+        text = prompt + "\n" + text
+        entities = [
+            {
+                "entity": entity["label"],
+                "word": entity["text"],
+                "start": entity["start"],
+                "end": entity["end"],
+                "score": 0,
+            }
+            for entity in gliner_model.predict_entities(
+                text, ["match"], flat_ner=not nested_ner, threshold=threshold
+            )
+        ]
+        merged_entities = merge_entities(entities)
+        return {"text": text, "entities": merged_entities}
+    except Exception as e:
+        return {"error": f"Information extraction failed: {e}"}
+
+def pipeline_fn(audio, prompt, threshold, nested_ner):
+    """
+    Combine transcription, translation, and information extraction in a single pipeline.
+    """
+    transcription = transcribe_audio(audio)
+    if "Error" in transcription:
+        return transcription, "", "", {}
+
+    translated_text = translate_text_to_english(transcription)
+    if "Error" in translated_text:
+        return transcription, translated_text, "", {}
+
+    info_extraction = extract_information(prompt, translated_text, threshold, nested_ner)
+    return transcription, translated_text, info_extraction
+
+# Gradio Interface
+with gr.Blocks(title="Audio Processing and Information Extraction") as interface:
+    gr.Markdown("## Audio Transcription, Translation, and Information Extraction")
+    
+    with gr.Row():
+        # Fixed: removed 'source' argument from gr.Audio
+        audio_input = gr.Audio(type="filepath", label="Upload Audio File")
+        prompt_input = gr.Textbox(label="Prompt for Information Extraction", placeholder="Enter your prompt here")
+    
+    with gr.Row():
+        threshold_slider = gr.Slider(0, 1, value=0.3, step=0.01, label="NER Threshold")
+        nested_ner_checkbox = gr.Checkbox(label="Enable Nested NER")
+    
+    with gr.Row():
+        transcription_output = gr.Textbox(label="Transcription (Urdu)", interactive=False)  # Corrected to interactive=False
+        translation_output = gr.Textbox(label="Translation (English)", interactive=False)  # Corrected to interactive=False
+    
+    with gr.Row():
+        extraction_output = gr.HighlightedText(label="Extracted Information")
+
+    process_button = gr.Button("Process Audio")
+
+    process_button.click(
+        fn=pipeline_fn,
+        inputs=[audio_input, prompt_input, threshold_slider, nested_ner_checkbox],
+        outputs=[transcription_output, translation_output, extraction_output],
+    )
+
+if __name__ == "__main__":
+    interface.launch()