Update app.py

app.py

@@ -1,47 +1,17 @@
-import sys
-import json
-from hugchat import hugchat
-from hugchat.login import Login
-import os
-import re
 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
 
-# Directory to store/load cookies
-cookie_path_dir = "./cookies/"
-cookie_file_path = os.path.join(cookie_path_dir, "cookies_snapshot.json")  # Default file name used by hugchat
-
-# Load pre-saved cookies
-try:
-    print("Attempting to load cookies from:", cookie_file_path)
-    if not os.path.exists(cookie_file_path):
-        # If cookies don't exist, attempt to generate them (for local testing; remove in Spaces)
-        EMAIL = os.environ.get("EMAIL", "[email protected]")  # Fallback for local testing
-        PASSWD = os.environ.get("PASSWORD", "e.AKsv$3Q4i4KcX")  # Fallback for local testing
-        os.makedirs(cookie_path_dir, exist_ok=True)
-        sign = Login(EMAIL, PASSWD)
-        cookies = sign.login(cookie_dir_path=cookie_path_dir, save_cookies=True)
-        print("Generated new cookies since none were found.")
-    else:
-        # Load existing cookies
-        with open(cookie_file_path, "r") as f:
-            cookies = json.load(f)  # Load the cookie dictionary
-        print("Cookies loaded from file.")
-    
-    chatbot = hugchat.ChatBot(cookies=cookies)  # Pass cookies directly
-    print("ChatBot initialized successfully.")
-except Exception as e:
-    print(f"Failed to initialize ChatBot: {str(e)}")
-    import traceback
-    traceback.print_exc()
-    sys.exit(1)
-
-# Model and device configuration for Whisper transcription
+# Model and device configuration for transcription
 MODEL_NAME = "openai/whisper-large-v3-turbo"
 device = 0 if torch.cuda.is_available() else "cpu"
 
+# Initialize Whisper pipeline
 pipe = pipeline(
     task="automatic-speech-recognition",
     model=MODEL_NAME,
@@ -49,139 +19,122 @@ pipe = pipeline(
     device=device,
 )
 
+# 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:
-        audio, sr = librosa.load(audio_path, sr=16000, mono=True)
-        transcription = pipe(audio, batch_size=8, generate_kwargs={"language": "urdu"})["text"]
-        return transcription
-    except Exception as e:
-        return f"Error processing audio: {e}"
-
-def extract_metadata(file_name):
-    base = file_name.split(".")[0]
-    parts = base.split("_")
-    if len(parts) >= 3:
-        return {
-            "agent_username": parts[0],
-            "location": parts[-2]
-        }
-    return {"agent_username": "Unknown", "location": "Unknown"}
+        # Log start time
+        start_time = datetime.now()
 
-def process_audio(audio, file_name):
-    urdu_text = transcribe_audio(audio)
-    if "Error" in urdu_text:
-        return json.dumps({"error": urdu_text})
-
-    metadata = extract_metadata(file_name)
-    location = metadata["location"]
+        # Ensure audio is mono and resampled to 16kHz
+        audio, sr = librosa.load(audio_path, sr=16000, mono=True)
 
-    english_text = chatbot.chat(
-        f"The following Urdu text is about crops and their diseases, but it may contain errors or misheard words due to audio transcription issues. Please use context to infer the most likely correct crop names and disease terms, and then translate the text to English:\n\n{urdu_text}"
-    ).wait_until_done()
+        # Perform transcription
+        transcription = pipe(audio, batch_size=8, generate_kwargs={"language": "urdu"})["text"]
 
-    extraction_prompt = f"""
-    Below is an English text about specific crops and possible diseases/pests:
+        # Log end time
+        end_time = datetime.now()
 
-    {english_text}
+        return transcription
 
-    Identify each specific Crop (like wheat, rice, cotton, etc.) mentioned and list any Diseases or Pests affecting that crop.
+    except Exception as e:
+        return f"Error processing audio: {e}"
 
-    - If a disease or pest is mentioned without specifying a particular crop, list it under "No crop:".
-    - If a crop is mentioned but no diseases or pests are specified for it, include it with an empty diseases list.
-    - Do not include general terms like "crops" as a specific crop name.
+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}"
 
-    Format your answer in this style (one entry at a time):
+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}"}
 
-    For specific crops with diseases:
-    1. CropName:
-    Diseases:
-    - DiseaseName
-    - AnotherDisease
+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, "", "", {}
 
-    For specific crops with no diseases:
-    2. NextCrop:
-    Diseases:
+    translated_text = translate_text_to_english(transcription)
+    if "Error" in translated_text:
+        return transcription, translated_text, "", {}
 
-    For standalone diseases:
-    3. No crop:
-    Diseases:
-    - StandaloneDisease
+    info_extraction = extract_information(prompt, translated_text, threshold, nested_ner)
+    return transcription, translated_text, info_extraction
 
-    No extra text, just the structured bullet list.
-    """
-    extraction_response = chatbot.chat(extraction_prompt).wait_until_done()
-
-    lines = extraction_response.splitlines()
-    crops_and_diseases = []
-    current_crop = None
-    current_diseases = []
-
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
-        match_crop = re.match(r'^(\d+)\.\s*(.+?):$', line)
-        if match_crop:
-            if current_crop is not None or current_diseases:
-                crops_and_diseases.append({
-                    "crop": current_crop,
-                    "diseases": current_diseases
-                })
-            crop_name = match_crop.group(2).strip()
-            if crop_name.lower() in ["no crop", "crops", "general crops"]:
-                current_crop = None
-            else:
-                current_crop = crop_name
-            current_diseases = []
-            continue
-        if line.lower().startswith("diseases:"):
-            continue
-        if line.startswith('-'):
-            disease_name = line.lstrip('-').strip()
-            if disease_name:
-                current_diseases.append(disease_name)
-
-    if current_crop is not None or current_diseases:
-        crops_and_diseases.append({
-            "crop": current_crop,
-            "diseases": current_diseases
-        })
-
-    temp_prompt = f"Give me weather of {location} in Celsius numeric only."
-    temperature_response = chatbot.chat(temp_prompt).wait_until_done()
-
-    temperature = None
-    temp_match = re.search(r'(\d+)', temperature_response)
-    if temp_match:
-        temperature = int(temp_match.group(1))
-
-    output = {
-        "urdu_text": urdu_text,
-        "english_text": english_text,
-        "crops_and_diseases": crops_and_diseases,
-        "temperature": temperature,
-        "location": location
-    }
-
-    return json.dumps(output)
-
-with gr.Blocks(title="Audio to Crop Disease API") as interface:
-    gr.Markdown("## Upload Audio to Get Urdu Transcription, English Translation, and Crop Diseases")
+# 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():
-        audio_input = gr.Audio(type="filepath", label="Upload Audio File (Urdu)")
-        file_name_input = gr.Textbox(label="File Name for Metadata (e.g., agent2_5_Multan_Pakistan.mp3)", placeholder="Enter file name")
+        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():
-        output_json = gr.JSON(label="Output (Urdu, English, Crops with Diseases)")
+        extraction_output = gr.HighlightedText(label="Extracted Information")
 
     process_button = gr.Button("Process Audio")
 
     process_button.click(
-        fn=process_audio,
-        inputs=[audio_input, file_name_input],
-        outputs=[output_json],
+        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()
+    interface.launch()