sagar007 commited on
Commit
1e235cc
·
verified ·
1 Parent(s): d0ce6f0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +119 -29
app.py CHANGED
@@ -1,26 +1,18 @@
1
  import gradio as gr
 
2
  import torch
3
- from transformers import AutoTokenizer, pipeline
4
 
5
  # Load the model and tokenizer
6
  model_name = "akjindal53244/Llama-3.1-Storm-8B"
7
  tokenizer = AutoTokenizer.from_pretrained(model_name)
8
- pipe = pipeline(
9
- "text-generation",
10
- model=model_name,
11
  torch_dtype=torch.bfloat16,
12
  device_map="auto"
13
  )
14
 
15
- # HTML content
16
- HTML_CONTENT = """
17
- <h1>Llama-3.1-Storm-8B Text Generation</h1>
18
- <p>Generate text using the powerful Llama-3.1-Storm-8B model. Enter a prompt and let the AI create!</p>
19
- <div class="llama-image">
20
- <img src="https://cdn-uploads.huggingface.co/production/uploads/64c75c1237333ccfef30a602/tmOlbERGKP7JSODa6T06J.jpeg" alt="Llama" style="width:200px; border-radius:10px;">
21
- </div>
22
- """
23
-
24
  def generate_text(prompt, max_length, temperature):
25
  messages = [
26
  {"role": "system", "content": "You are a helpful assistant."},
@@ -28,8 +20,10 @@ def generate_text(prompt, max_length, temperature):
28
  ]
29
  formatted_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
30
 
31
- outputs = pipe(
32
- formatted_prompt,
 
 
33
  max_new_tokens=max_length,
34
  do_sample=True,
35
  temperature=temperature,
@@ -37,20 +31,116 @@ def generate_text(prompt, max_length, temperature):
37
  top_p=0.95,
38
  )
39
 
40
- return outputs[0]['generated_text'][len(formatted_prompt):]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
 
42
- with gr.Blocks() as demo:
43
- gr.HTML(HTML_CONTENT)
44
- with gr.Row():
45
- with gr.Column(scale=2):
46
- prompt = gr.Textbox(label="Prompt", lines=5)
47
- max_length = gr.Slider(minimum=1, maximum=500, value=128, step=1, label="Max Length")
48
- temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
49
- submit_button = gr.Button("Generate")
50
- with gr.Column(scale=2):
51
- output = gr.Textbox(label="Generated Text", lines=10)
52
 
53
- submit_button.click(generate_text, inputs=[prompt, max_length, temperature], outputs=[output])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
54
 
55
- if __name__ == "__main__":
56
- demo.launch()
 
1
  import gradio as gr
2
+ import spaces
3
  import torch
4
+ from transformers import AutoTokenizer, AutoModelForCausalLM
5
 
6
  # Load the model and tokenizer
7
  model_name = "akjindal53244/Llama-3.1-Storm-8B"
8
  tokenizer = AutoTokenizer.from_pretrained(model_name)
9
+ model = AutoModelForCausalLM.from_pretrained(
10
+ model_name,
 
11
  torch_dtype=torch.bfloat16,
12
  device_map="auto"
13
  )
14
 
15
+ @spaces.GPU(duration=120)
 
 
 
 
 
 
 
 
16
  def generate_text(prompt, max_length, temperature):
17
  messages = [
18
  {"role": "system", "content": "You are a helpful assistant."},
 
20
  ]
21
  formatted_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
22
 
23
+ inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
24
+
25
+ outputs = model.generate(
26
+ **inputs,
27
  max_new_tokens=max_length,
28
  do_sample=True,
29
  temperature=temperature,
 
31
  top_p=0.95,
32
  )
33
 
34
+ return tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
35
+
36
+ # Custom CSS
37
+ css = """
38
+ body {
39
+ background: linear-gradient(135deg, #f5f7fa, #c3cfe2);
40
+ font-family: Arial, sans-serif;
41
+ }
42
+ #custom-header {
43
+ text-align: center;
44
+ background: rgba(255, 255, 255, 0.8);
45
+ padding: 20px;
46
+ border-radius: 10px;
47
+ box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
48
+ position: relative;
49
+ max-width: 800px;
50
+ margin: 20px auto;
51
+ }
52
+ #custom-header h1 {
53
+ color: #4A90E2;
54
+ font-size: 2em;
55
+ margin-bottom: 10px;
56
+ }
57
+ .llama-image {
58
+ position: relative;
59
+ transition: transform 0.3s;
60
+ display: inline-block;
61
+ margin-top: 20px;
62
+ }
63
+ .llama-image:hover {
64
+ transform: scale(1.05);
65
+ }
66
+ .llama-image img {
67
+ width: 200px;
68
+ border-radius: 10px;
69
+ box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
70
+ }
71
+ .llama-description {
72
+ position: absolute;
73
+ bottom: -30px;
74
+ left: 50%;
75
+ transform: translateX(-50%);
76
+ background-color: #4A90E2;
77
+ color: white;
78
+ padding: 5px 10px;
79
+ border-radius: 5px;
80
+ opacity: 0;
81
+ transition: opacity 0.3s;
82
+ white-space: nowrap;
83
+ }
84
+ .llama-image:hover .llama-description {
85
+ opacity: 1;
86
+ }
87
+ .artifact {
88
+ position: absolute;
89
+ background: rgba(74, 144, 226, 0.1);
90
+ border-radius: 50%;
91
+ }
92
+ .artifact.large {
93
+ width: 300px;
94
+ height: 300px;
95
+ top: -50px;
96
+ left: -150px;
97
+ }
98
+ .artifact.medium {
99
+ width: 200px;
100
+ height: 200px;
101
+ bottom: -50px;
102
+ right: -100px;
103
+ }
104
+ .artifact.small {
105
+ width: 100px;
106
+ height: 100px;
107
+ top: 50%;
108
+ left: 50%;
109
+ transform: translate(-50%, -50%);
110
+ }
111
+ """
112
 
113
+ # Custom HTML
114
+ custom_html = """
115
+ <div id="custom-header">
116
+ <div class="artifact large"></div>
117
+ <div class="artifact medium"></div>
118
+ <div class="artifact small"></div>
 
 
 
 
119
 
120
+ <h1>Llama-3.1-Storm-8B Text Generation</h1>
121
+ <p>Generate text using the powerful Llama-3.1-Storm-8B model. Enter a prompt and let the AI create!</p>
122
+
123
+ <div class="llama-image">
124
+ <img src="https://cdn-uploads.huggingface.co/production/uploads/64c75c1237333ccfef30a602/tmOlbERGKP7JSODa6T06J.jpeg" alt="Llama">
125
+ <div class="llama-description">Llama-3.1-Storm-8B Model</div>
126
+ </div>
127
+ </div>
128
+ """
129
+
130
+ # Create Gradio interface
131
+ iface = gr.Interface(
132
+ fn=generate_text,
133
+ inputs=[
134
+ gr.Textbox(lines=5, label="Prompt"),
135
+ gr.Slider(minimum=1, maximum=500, value=128, step=1, label="Max Length"),
136
+ gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
137
+ ],
138
+ outputs=gr.Textbox(lines=10, label="Generated Text"),
139
+ title="Llama-3.1-Storm-8B Text Generation",
140
+ description="Enter a prompt to generate text using the Llama-3.1-Storm-8B model.",
141
+ css=css,
142
+ article=custom_html
143
+ )
144
 
145
+ # Launch the app
146
+ iface.launch()