app.py

import subprocess 

subprocess.run(
    'pip install flash-attn --no-build-isolation',
    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
    shell=True
)
import os
import re
import time
import torch
import spaces
import gradio as gr
from threading import Thread
from transformers import (
    AutoModelForCausalLM, 
    AutoTokenizer, 
    BitsAndBytesConfig, 
    TextIteratorStreamer
)

# Configuration Constants
MODEL_ID = "CohereForAI/aya-expanse-8b"
DEFAULT_SYSTEM_PROMPT = """ 
أنت مترجم ثنائي اللغة متخصص في الترجمة بين العربية والإنجليزية. هدفك هو تقديم ترجمات دقيقة، ملائمة للسياق، ومتسقة من الناحية الأسلوبية، مع الالتزام بالإرشادات التالية:

أسلوب الكتابة:

1. الدقة النحوية: احرص دائمًا على أن تكون الترجمة صحيحة نحويًا.


2. الملاءمة السياقية: قم بتكييف الترجمة مع السياق والجمهور المستهدف:

استخدم جملًا موجزة وواضحة في المحتوى الطبي والقانوني.

اعتمد لغة جذابة وخيالية في المواد التسويقية.

حافظ على الرسمية والبلاغة في المحتوى القانوني.



3. البنية السليمة: احترم تراكيب الجمل العربية وتجنب النسخ المباشر لقواعد لغة المصدر. فضل استخدام "الجملة الفعلية" إلا إذا كانت "الجملة الاسمية" أكثر ملاءمة (مثلًا في العناوين أو التنويهات).



الخيارات الأسلوبية:

استخدم التشكيل فقط عند الضرورة لتوضيح المعنى.

تعامل مع الأسماء والأحرف المختصرة حسب السياق:

قم بتعريب الأسماء وأسماء الأدوية إلا إذا وُجدت لها مرادفات عربية.

ترجم أسماء البرامج والإدارات والهيئات عندما يكون ذلك مفيدًا.


استخدم الأرقام العربية وتأكد من التعامل الصحيح مع الوحدات والعناوين والمراجع.


علامات الترقيم:

طبق قواعد الترقيم العربية لضمان وضوح النص.

استخدم الفاصلة العربية (،) والفاصلة المنقوطة (؛) حسب الأعراف.

تجنب الإفراط في استخدام علامات الاقتباس واحرص على وضع النقطتين (:) بشكل منطقي.


الأخطاء الشائعة التي يجب تجنبها:

تجنب ترجمة "is" بـ "يُعتبر" إلا إذا كان ذلك مناسبًا للسياق.

استخدم حروف الجر والعطف بشكل صحيح لضمان تدفق الجمل بسلاسة.

قلل من التكرار واستخدم الضمائر حيثما أمكن.

تجنب الإفراط في استخدام تراكيب مثل "(قام + الفعل)" و"الخاص بـ."


المصطلحات المحددة:

في الترجمات القانونية، حافظ على النبرة الرسمية وتأكد من دقة المصطلحات.

في الترجمات الطبية، بسّط المصطلحات التقنية للجمهور العام مع الاحتفاظ بالتعقيد عند مخاطبة المتخصصين.

في الترجمات التسويقية، ركّز على الإبداع بدلًا من الترجمة الحرفية، مع الحفاظ على الرسالة الأساسية.


إرشادات التنسيق:

التزم بمعايير الطباعة العربية بشكل متسق.

حافظ على تنسيق البيانات المهمة (مثل التواريخ، والقياسات، والاستشهادات القانونية).


عند الشك:

قدم الأولوية للوضوح، والتناسق، والملاءمة مع احتياجات الجمهور المستهدف. قم دائمًا بموازنة التعليمات الخاصة بالمشروع مع هذه الإرشادات، مع إعطاء الأولوية لمتطلبات العميل عند وجود أي تعارض.

أنت مترجم ثنائي اللغة متخصص في الترجمة بين العربية والإنجليزية. هدفك هو تقديم ترجمات دقيقة، ملائمة للسياق، ومتسقة من الناحية الأسلوبية، مع الالتزام بالإرشادات التالية:

أسلوب الكتابة:

1. الدقة النحوية: احرص دائمًا على أن تكون الترجمة صحيحة نحويًا.


2. الملاءمة السياقية: قم بتكييف الترجمة مع السياق والجمهور المستهدف:

استخدم جملًا موجزة وواضحة في المحتوى الطبي والقانوني.

اعتمد لغة جذابة وخيالية في المواد التسويقية.

حافظ على الرسمية والبلاغة في المحتوى القانوني.



3. البنية السليمة: احترم تراكيب الجمل العربية وتجنب النسخ المباشر لقواعد لغة المصدر. فضل استخدام "الجملة الفعلية" إلا إذا كانت "الجملة الاسمية" أكثر ملاءمة (مثلًا في العناوين أو التنويهات).



الخيارات الأسلوبية:

استخدم التشكيل فقط عند الضرورة لتوضيح المعنى.

تعامل مع الأسماء والأحرف المختصرة حسب السياق:

قم بتعريب الأسماء وأسماء الأدوية إلا إذا وُجدت لها مرادفات عربية.

ترجم أسماء البرامج والإدارات والهيئات عندما يكون ذلك مفيدًا.


استخدم الأرقام العربية وتأكد من التعامل الصحيح مع الوحدات والعناوين والمراجع.


علامات الترقيم:

طبق قواعد الترقيم العربية لضمان وضوح النص.

استخدم الفاصلة العربية (،) والفاصلة المنقوطة (؛) حسب الأعراف.

تجنب الإفراط في استخدام علامات الاقتباس واحرص على وضع النقطتين (:) بشكل منطقي.


الأخطاء الشائعة التي يجب تجنبها:

تجنب ترجمة "is" بـ "يُعتبر" إلا إذا كان ذلك مناسبًا للسياق.

استخدم حروف الجر والعطف بشكل صحيح لضمان تدفق الجمل بسلاسة.

قلل من التكرار واستخدم الضمائر حيثما أمكن.

تجنب الإفراط في استخدام تراكيب مثل "(قام + الفعل)" و"الخاص بـ."


المصطلحات المحددة:

في الترجمات القانونية، حافظ على النبرة الرسمية وتأكد من دقة المصطلحات.

في الترجمات الطبية، بسّط المصطلحات التقنية للجمهور العام مع الاحتفاظ بالتعقيد عند مخاطبة المتخصصين.

في الترجمات التسويقية، ركّز على الإبداع بدلًا من الترجمة الحرفية، مع الحفاظ على الرسالة الأساسية.


إرشادات التنسيق:

التزم بمعايير الطباعة العربية بشكل متسق.

حافظ على تنسيق البيانات المهمة (مثل التواريخ، والقياسات، والاستشهادات القانونية).


عند الشك:

قدم الأولوية للوضوح، والتناسق، والملاءمة مع احتياجات الجمهور المستهدف. قم دائمًا بموازنة التعليمات الخاصة بالمشروع مع هذه الإرشادات، مع إعطاء الأولوية لمتطلبات العميل عند وجود أي تعارض.


"""
# UI Configuration
TITLE = "<h1><center>Mawared T Assistant</center></h1>"
PLACEHOLDER = "Ask me anything! I'll think through it step by step."

CSS = """
.duplicate-button {
    margin: auto !important;
    color: white !important;
    background: black !important;
    border-radius: 100vh !important;
}
h3 {
    text-align: center;
}
.message-wrap {
    overflow-x: auto;
}
.message-wrap p {
    margin-bottom: 1em;
}
.message-wrap pre {
    background-color: #f6f8fa;
    border-radius: 3px;
    padding: 16px;
    overflow-x: auto;
}
.message-wrap code {
    background-color: rgba(175,184,193,0.2);
    border-radius: 3px;
    padding: 0.2em 0.4em;
    font-family: monospace;
}
.custom-tag {
    color: #0066cc;
    font-weight: bold;
}
.chat-area {
    height: 500px !important;
    overflow-y: auto !important;
}
"""

def initialize_model():
    """Initialize the model with appropriate configurations"""
    quantization_config = BitsAndBytesConfig(
        load_in_8bit=True,
        bnb_8bit_compute_dtype=torch.bfloat16,
        bnb_8bit_use_double_quant=True
    )

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        device_map="cuda",
        attn_implementation="flash_attention_2",
        #quantization_config=quantization_config

    )

    return model, tokenizer

def format_text(text):
    """Format text with proper spacing and tag highlighting (but keep tags visible)"""
    tag_patterns = [
        (r'<Thinking>', '\n<Thinking>\n'),
        (r'</Thinking>', '\n</Thinking>\n'),
        (r'<Critique>', '\n<Critique>\n'),
        (r'</Critique>', '\n</Critique>\n'),
        (r'<Revising>', '\n<Revising>\n'),
        (r'</Revising>', '\n</Revising>\n'),
        (r'<Final>', '\n<Final>\n'),
        (r'</Final>', '\n</Final>\n')
    ]
    
    formatted = text
    for pattern, replacement in tag_patterns:
        formatted = re.sub(pattern, replacement, formatted)
    
    formatted = '\n'.join(line for line in formatted.split('\n') if line.strip())
    
    return formatted

def format_chat_history(history):
    """Format chat history for display, keeping tags visible"""
    formatted = []
    for user_msg, assistant_msg in history:
        formatted.append(f"User: {user_msg}")
        if assistant_msg:
            formatted.append(f"Assistant: {assistant_msg}")
    return "\n\n".join(formatted)
    
def create_examples():
    """Create example queries for the UI"""
    return [
        "Explain the concept of artificial intelligence.",
        "How does photosynthesis work?",
        "What are the main causes of climate change?",
        "Describe the process of protein synthesis.",
        "What are the key features of a democratic government?",
        "Explain the theory of relativity.",
        "How do vaccines work to prevent diseases?",
        "What are the major events of World War II?",
        "Describe the structure of a human cell.",
        "What is the role of DNA in genetics?"
    ]

@spaces.GPU(duration=660)
def chat_response(
    message: str,
    history: list,
    chat_display: str,
    system_prompt: str,
    temperature: float = 0.2,
    max_new_tokens: int = 8192,
    top_p: float = 0.8,
    top_k: int = 40,
    penalty: float = 1.2,
):
    """Generate chat responses, keeping tags visible in the output"""
    conversation = [
        {"role": "system", "content": system_prompt}
    ]
    
    for prompt, answer in history:
        conversation.extend([
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": answer}
        ])
    
    conversation.append({"role": "user", "content": message})
    
    input_ids = tokenizer.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)
    
    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=60.0,
        skip_prompt=True,
        skip_special_tokens=True
    )
    
    generate_kwargs = dict(
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=False if temperature == 0 else True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        repetition_penalty=penalty,
        streamer=streamer,
    )
    
    buffer = ""
    
    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
        
        history = history + [[message, ""]]
        
        for new_text in streamer:
            buffer += new_text
            formatted_buffer = format_text(buffer)
            history[-1][1] = formatted_buffer
            chat_display = format_chat_history(history)
            
            yield history, chat_display

def process_example(example: str) -> tuple:
    """Process example query and return empty history and updated display"""
    return [], f"User: {example}\n\n"

def main():
    """Main function to set up and launch the Gradio interface"""
    global model, tokenizer
    model, tokenizer = initialize_model()
    
    with gr.Blocks(css=CSS, theme="soft") as demo:
        gr.HTML(TITLE)
        gr.DuplicateButton(
            value="Duplicate Space for private use",
            elem_classes="duplicate-button"
        )
        
        with gr.Row():
            with gr.Column():
                chat_history = gr.State([])
                chat_display = gr.TextArea(
                    value="",
                    label="Chat History",
                    interactive=False,
                    elem_classes=["chat-area"],
                )
                
                message = gr.TextArea(
                    placeholder=PLACEHOLDER,
                    label="Your message",
                    lines=3
                )
                
                with gr.Row():
                    submit = gr.Button("Send")
                    clear = gr.Button("Clear")
                
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    system_prompt = gr.TextArea(
                        value=DEFAULT_SYSTEM_PROMPT,
                        label="System Prompt",
                        lines=5,
                    )
                    temperature = gr.Slider(
                        minimum=0,
                        maximum=1,
                        step=0.1,
                        value=0.2,
                        label="Temperature",
                    )
                    max_tokens = gr.Slider(
                        minimum=128,
                        maximum=32000,
                        step=128,
                        value=8192,
                        label="Max Tokens",
                    )
                    top_p = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        step=0.1,
                        value=0.8,
                        label="Top-p",
                    )
                    top_k = gr.Slider(
                        minimum=1,
                        maximum=100,
                        step=1,
                        value=40,
                        label="Top-k",
                    )
                    penalty = gr.Slider(
                        minimum=1.0,
                        maximum=2.0,
                        step=0.1,
                        value=1.2,
                        label="Repetition Penalty",
                    )
                
                examples = gr.Examples(
                    examples=create_examples(),
                    inputs=[message],
                    outputs=[chat_history, chat_display],
                    fn=process_example,
                    cache_examples=False,
                )
        
        # Set up event handlers
        submit_click = submit.click(
            chat_response,
            inputs=[
                message,
                chat_history,
                chat_display,
                system_prompt,
                temperature,
                max_tokens,
                top_p,
                top_k,
                penalty,
            ],
            outputs=[chat_history, chat_display],
            show_progress=True,
        )
        
        message.submit(
            chat_response,
            inputs=[
                message,
                chat_history,
                chat_display,
                system_prompt,
                temperature,
                max_tokens,
                top_p,
                top_k,
                penalty,
            ],
            outputs=[chat_history, chat_display],
            show_progress=True,
        )
        
        clear.click(
            lambda: ([], ""),
            outputs=[chat_history, chat_display],
            show_progress=True,
        )
        
        submit_click.then(lambda: "", outputs=message)
        message.submit(lambda: "", outputs=message)
    
    return demo

if __name__ == "__main__":
    demo = main()
    demo.launch()