geeaiml commited on
Commit
d28df4e
·
verified ·
1 Parent(s): 2fb48d1

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +105 -0
app.py ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
3
+ import torch
4
+ import spaces
5
+ import os
6
+
7
+ import gc
8
+ import torch
9
+
10
+ # Create the necessary directories
11
+ os.makedirs('.gradio/cached_examples/17', exist_ok=True)
12
+
13
+ def get_model_name(language):
14
+ """Map language choice to the corresponding model."""
15
+ model_mapping = {
16
+ "English": "microsoft/Phi-3-mini-4k-instruct",
17
+ "Arabic": "ALLaM-AI/ALLaM-7B-Instruct-preview"
18
+ }
19
+ return model_mapping.get(language, "ALLaM-AI/ALLaM-7B-Instruct-preview") # Default to Arabic model
20
+
21
+ def load_model(model_name):
22
+ device = "cuda" if torch.cuda.is_available() else "cpu"
23
+ model = AutoModelForCausalLM.from_pretrained(
24
+ model_name,
25
+ device_map=device,
26
+ torch_dtype="auto",
27
+ trust_remote_code=True,
28
+ )
29
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
30
+ generator = pipeline(
31
+ "text-generation",
32
+ model=model,
33
+ tokenizer=tokenizer,
34
+ return_full_text=False,
35
+ max_new_tokens=500,
36
+ do_sample=True, # Enable sampling for more creative outputs
37
+ top_k=50, # Control diversity
38
+ top_p=0.95 # Control diversity
39
+ )
40
+ del model
41
+ del tokenizer
42
+
43
+ return generator
44
+
45
+
46
+ @spaces.GPU
47
+ def generate_kids_story(character, setting, language):
48
+ model_name = get_model_name(language)
49
+ generator = load_model(model_name)
50
+
51
+ # Define prompt for the AI model
52
+ if language == "English":
53
+ prompt = (f"Write a short story for kids about a character named {character} who goes on an adventure in {setting}. "
54
+ "Make it fun, engaging, and suitable for children.")
55
+ else:
56
+ prompt = (f"اكتب قصة قصيرة للأطفال عن شخصية اسمها {character} التي تذهب في مغامرة في {setting}. "
57
+ "اجعلها ممتعة وجذابة ومناسبة للأطفال.")
58
+
59
+ messages = [{"role": "user", "content": prompt}]
60
+ output = generator(messages)
61
+
62
+ # Delete model and associated objects
63
+ del generator
64
+ # Run garbage collection
65
+ gc.collect ()
66
+ # Empty CUDA cache
67
+ torch.cuda.empty_cache()
68
+
69
+ return output[0]["generated_text"]
70
+
71
+
72
+ css_style = """
73
+ body {
74
+ background-image: url('https://cdna.artstation.com/p/assets/images/images/074/776/904/large/pietro-chiovaro-r1-castle-chp.jpg?1712916847');
75
+ background-size: cover;
76
+ background-position: center;
77
+ color: #fff; /* General text color */
78
+ font-family: 'Arial', sans-serif;
79
+ }"""
80
+ # Create Gradio interface
81
+ demo = gr.Interface(
82
+ fn=generate_kids_story,
83
+ inputs=[
84
+ gr.Textbox(placeholder="Enter a character name (e.g., Benny the Bunny)...", label="Character Name"),
85
+ gr.Textbox(placeholder="Enter a setting (e.g., a magical forest)...", label="Setting"),
86
+ gr.Dropdown(
87
+ choices=["English", "Arabic"],
88
+ label="Choose Language",
89
+ value="English" # Default to English
90
+ )
91
+ ],
92
+ outputs=gr.Textbox(label="Kids' Story"),
93
+ title="📖 AI Kids' Story Generator - English & Arabic 📖",
94
+ description="Enter a character name and a setting, and AI will generate a fun short story for kids in English or Arabic.",
95
+ examples=[
96
+ ["Benny the Bunny", "a magical forest", "English"],
97
+ ["علي البطل", "غابة سحرية", "Arabic"],
98
+ ["Lila the Ladybug", "a garden full of flowers", "English"],
99
+ ["ليلى الجنية", "حديقة مليئة بالأزهار", "Arabic"]
100
+ ],
101
+ css= css_style,
102
+ )
103
+
104
+ # Launch the Gradio app
105
+ demo.launch()