app.py

import gradio as gr
import time
import random

# This import will fail unless your environment includes the internal "spaces" module
try:
    import spaces
except ImportError:
    raise ImportError(
        "Could not import 'spaces'. The @spaces.GPU decorator is likely not available in your environment.\n"
        "To use GPU, create an empty 'gpu' file in your Space or enable GPU in Space Settings."
    )

from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# CSS for custom styling with RTL support
custom_css = """
.gradio-container {
    background: linear-gradient(to bottom right, #1a1a2e, #16213e);
}
.container {
    border-radius: 15px;
    box-shadow: 0 10px 25px rgba(0, 0, 0, 0.3);
    padding: 20px;
    background-color: #0f3460;
    margin-bottom: 20px;
}
footer {
    color: #e94560 !important;
}
.story-output textarea {
    background-color: #16213e !important;
    border-radius: 10px !important;
    padding: 15px !important;
    font-family: 'Arial', sans-serif !important;
    line-height: 1.8 !important;
    color: #ffffff !important;
    direction: rtl !important;  /* Right-to-left for Arabic */
    text-align: right !important;
    font-size: 16px !important;
}
.example-btn button {
    background-color: #e94560 !important;
    color: white !important;
}
.info-text {
    font-size: 0.9em;
    color: #cccccc;
}
.title {
    text-align: center;
    color: #e94560 !important;
    font-size: 2.5em !important;
    margin-bottom: 0;
}
.subtitle {
    text-align: center;
    color: #cccccc !important;
    font-size: 1.2em !important;
    margin-top: 5px;
}
/* RTL support for prompt input */
.rtl-text textarea {
    direction: rtl !important;
    text-align: right !important;
    font-size: 16px !important;
}
"""

# Load the model and tokenizer
print("Loading model and tokenizer...")
model_name = "ALLaM-AI/ALLaM-7B-Instruct-preview"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Move model to GPU immediately
model.to("cuda")

# Create a pipeline
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)

# Example prompts for the user to try (in Arabic)
example_prompts = [
    "في عالم حيث يمكن مشاركة الأحلام، يكتشف محقق شاب أن قاتلًا متسلسلًا يستخدم الأحلام لاصطياد ضحاياه...",
    "يكتشف شقيقان خريطة قديمة في علية جدتهما تقودهما إلى عالم سحري مخفي في الغابة خلف منزلهما...",
    "طباخ لديه القدرة على تذوق المشاعر في الطعام يُحضّر وجبة لزبون غامض يبدو أنه بلا مشاعر على الإطلاق...",
    "بعد حدث عالمي تسبب في توقف جميع التكنولوجيا عن العمل، تكتشف مجموعة من المراهقين أن لديهم قدرات جديدة غريبة...",
    "يصل مسافر عبر الزمن إلى يومنا الحاضر مع تحذير عاجل، لكن لا أحد يصدقه لأن...",
]

# Story genres for the dropdown (in Arabic)
story_genres = [
    "مغامرة",
    "خيال",
    "خيال علمي",
    "غموض",
    "رومانسية",
    "رعب",
    "كوميديا",
    "دراما",
    "حكاية خيالية",
    "مستقبل قاتم"
]

# Decorate the function that needs GPU with @spaces.GPU
@spaces.GPU
def generate_story(prompt, max_length=300, temperature=0.7, top_p=0.9, genre="مغامرة", creativity_level="متوسط"):
    """
    Generate a story based on the user's prompt and specified parameters.
    GPU is allocated only while this function is running in ZeroGPU environments.
    """
    # Enhance the prompt with the genre and creativity instructions (in Arabic)
    creativity_map = {
        "تقليدي": "اكتب قصة تقليدية ومباشرة من نوع",
        "متوسط": "اكتب قصة إبداعية وجذابة من نوع",
        "إبداعي": "اكتب قصة خيالية وفريدة من نوع"
    }
    
    enhanced_prompt = f"{creativity_map.get(creativity_level, 'اكتب قصة جذابة من نوع')} {genre}: {prompt}"
    
    # Generate the story
    response = generator(
        enhanced_prompt,
        max_length=max_length,
        temperature=float(temperature),
        top_p=float(top_p),
        do_sample=True
    )
    
    generated_text = response[0]["generated_text"]
    
    # Format the story to start from where the prompt ends
    if generated_text.startswith(enhanced_prompt):
        formatted_story = generated_text[len(enhanced_prompt):].strip()
    else:
        formatted_story = generated_text
    
    # Return the generated story
    return formatted_story

# Function to get a random example prompt
def get_random_example():
    return random.choice(example_prompts)

# Function to update the UI based on creativity level
def update_creativity_params(creativity):
    if creativity == "تقليدي":
        return 0.5, 0.85
    elif creativity == "متوسط":
        return 0.7, 0.9
    elif creativity == "إبداعي":
        return 0.9, 0.95
    return 0.7, 0.9

# Main function for Gradio interface
def create_interface():
    with gr.Blocks(css=custom_css) as demo:
        # Title and subtitle (in Arabic)
        gr.HTML("<h1 class='title'>✨ مولد القصص الذكي ✨</h1>")
        gr.HTML("<p class='subtitle'>أطلق العنان لخيالك مع سرد القصص بمساعدة الذكاء الاصطناعي</p>")
        
        # Input section
        with gr.Row():
            with gr.Column():
                gr.HTML("<div class='container'>")
                prompt = gr.Textbox(
                    placeholder="أدخل بداية قصتك هنا...",
                    label="بداية القصة",
                    lines=5,
                    elem_classes=["rtl-text"]  # Add RTL class
                )
                
                with gr.Row():
                    genre = gr.Dropdown(
                        choices=story_genres,
                        value="مغامرة",
                        label="نوع القصة"
                    )
                    creativity = gr.Radio(
                        choices=["تقليدي", "متوسط", "إبداعي"],
                        value="متوسط",
                        label="مستوى الإبداع"
                    )
                
                with gr.Accordion("إعدادات متقدمة", open=False):
                    with gr.Row():
                        temperature = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            step=0.05,
                            value=0.7,
                            label="درجة الحرارة"
                        )
                        top_p = gr.Slider(
                            minimum=0.5,
                            maximum=1.0,
                            step=0.05,
                            value=0.9,
                            label="Top-P"
                        )
                    
                    max_length = gr.Slider(
                        minimum=100,
                        maximum=1500,
                        step=50,
                        value=500,
                        label="الحد الأقصى لطول القصة"
                    )
                
                gr.HTML("<div class='info-text' dir='rtl'><p>✨ <strong>نصائح لبدايات قصص رائعة:</strong> كن محددًا، أضف شخصيات وإعدادات وصراعات.</p></div>")
                
                with gr.Row():
                    submit_btn = gr.Button("توليد القصة", variant="primary")
                    with gr.Column(elem_classes=["example-btn"]):
                        example_btn = gr.Button("جرب مثالًا")
                gr.HTML("</div>")  # Close container div
                
        # Output section
        with gr.Row():
            with gr.Column():
                gr.HTML("<div class='container'>")
                output = gr.Textbox(
                    label="قصتك",
                    lines=12,
                    elem_classes=["story-output"]
                )
                
                gr.HTML("<div class='info-text' style='text-align: center; margin-top: 10px;' dir='rtl'><p>💡 كل قصة يتم إنشاؤها بشكل فريد وقد تختلف في الأسلوب والمحتوى.</p></div>")
                gr.HTML("</div>")  # Close container div
        
        # Set up the submit action
        submit_btn.click(
            fn=generate_story,
            inputs=[prompt, max_length, temperature, top_p, genre, creativity],
            outputs=[output]
        )
        
        # Get random example when button is clicked
        example_btn.click(
            fn=get_random_example,
            outputs=[prompt]
        )
        
        # Update parameters based on creativity level
        creativity.change(
            fn=update_creativity_params,
            inputs=[creativity],
            outputs=[temperature, top_p]
        )
                
        gr.HTML("<div style='text-align: center; margin-top: 20px; color: #cccccc;' dir='rtl'><p>مدعوم بواسطة ALLaM-7B-Instruct • تم إنشاؤه باستخدام Gradio</p></div>")
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_interface()
    demo.launch()