app.py

import gradio as gr
from faster_whisper import WhisperModel
from groq import Groq
import tempfile
from gtts import gTTS
import fitz  # PyMuPDF

# Load Faster Whisper model (CPU optimized)
whisper_model = WhisperModel("base", compute_type="int8")

# Set Groq API Key
groq_api_key = "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"
client = Groq(api_key=groq_api_key)

# Urdu Legal Chat
def urdu_chat(question):
    response = client.chat.completions.create(
        model="llama3-8b-8192",  # Smaller model for faster response
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant that answers legal questions in Urdu simply."},
            {"role": "user", "content": question}
        ]
    )
    return response.choices[0].message.content

# Voice-based Urdu Chat
def urdu_chat_voice(audio_file_path):
    segments, _ = whisper_model.transcribe(audio_file_path, language="ur")
    question = " ".join([seg.text for seg in segments])

    response = client.chat.completions.create(
        model="llama3-8b-8192",
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant that answers legal questions in Urdu simply."},
            {"role": "user", "content": question}
        ]
    )
    answer = response.choices[0].message.content

    tts = gTTS(answer, lang='ur')
    audio_path = tempfile.mktemp(suffix=".mp3")
    tts.save(audio_path)

    return question, answer, audio_path

# PDF Legal Analysis
def analyze_pdf(pdf_file):
    text = ""
    with fitz.open(pdf_file.name) as doc:
        for page in doc:
            text += page.get_text()
    prompt = f"""تم ایک ماہر اردو قانونی تجزیہ نگار ہو۔ درج ذیل دستاویز کا جائزہ لے کر سادہ اردو میں اپنی رائے دو:
{text}"""
    response = client.chat.completions.create(
        model="llama3-8b-8192",
        messages=[
            {"role": "system", "content": "You are an expert Urdu legal analyst."},
            {"role": "user", "content": prompt}
        ]
    )
    return response.choices[0].message.content

# Gradio UI
with gr.Blocks(title="⚖ Urdu Legal Assistant") as demo:
    gr.Markdown("## 🧑‍⚖ اردو میں قانونی سوالات پوچھیں (ٹیکسٹ، آڈیو، یا دستاویز کے ذریعے)")

    with gr.Tab("✍ ٹیکسٹ چیٹ"):
        with gr.Row():
            text_input = gr.Textbox(lines=2, placeholder="اپنا قانونی سوال یہاں لکھیں...")
            text_output = gr.Textbox(label="جواب")
        text_btn = gr.Button("جواب حاصل کریں")
        text_btn.click(fn=urdu_chat, inputs=text_input, outputs=text_output)

    with gr.Tab("🎙 آواز سے پوچھیں"):
        with gr.Row():
            audio_input = gr.Audio(type="filepath", label="🔊 اپنی آواز میں سوال ریکارڈ کریں یا اپلوڈ کریں")
        with gr.Row():
            transcribed = gr.Textbox(label="📝 سوال (آڈیو سے متن)")
            ai_answer = gr.Textbox(label="🤖 AI جواب")
        with gr.Row():
            audio_response = gr.Audio(type="filepath", label="🔊 آڈیو میں جواب سنیں")
        voice_btn = gr.Button("جواب سنیں")
        voice_btn.click(fn=urdu_chat_voice, inputs=audio_input, outputs=[transcribed, ai_answer, audio_response])

    with gr.Tab("📄 PDF دستاویز"):
        with gr.Row():
            pdf_input = gr.File(label="قانونی PDF دستاویز اپلوڈ کریں", file_types=[".pdf"])
        with gr.Row():
            pdf_output = gr.Textbox(label="⚖ قانونی تجزیہ")
        pdf_btn = gr.Button("تجزیہ کریں")
        pdf_btn.click(fn=analyze_pdf, inputs=pdf_input, outputs=pdf_output)

demo.launch()