Update app.py

app.py

@@ -1,26 +1,18 @@
 import gradio as gr
+import spaces
 import torch
-from transformers import AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the model and tokenizer
 model_name = "akjindal53244/Llama-3.1-Storm-8B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-pipe = pipeline(
-    "text-generation",
-    model=model_name,
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
     torch_dtype=torch.bfloat16,
     device_map="auto"
 )
 
-# HTML content
-HTML_CONTENT = """
-<h1>Llama-3.1-Storm-8B Text Generation</h1>
-<p>Generate text using the powerful Llama-3.1-Storm-8B model. Enter a prompt and let the AI create!</p>
-<div class="llama-image">
-    <img src="https://cdn-uploads.huggingface.co/production/uploads/64c75c1237333ccfef30a602/tmOlbERGKP7JSODa6T06J.jpeg" alt="Llama" style="width:200px; border-radius:10px;">
-</div>
-"""
-
+@spaces.GPU(duration=120)
 def generate_text(prompt, max_length, temperature):
     messages = [
         {"role": "system", "content": "You are a helpful assistant."},
@@ -28,8 +20,10 @@ def generate_text(prompt, max_length, temperature):
     ]
     formatted_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
     
-    outputs = pipe(
-        formatted_prompt,
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+    
+    outputs = model.generate(
+        **inputs,
         max_new_tokens=max_length,
         do_sample=True,
         temperature=temperature,
@@ -37,20 +31,116 @@ def generate_text(prompt, max_length, temperature):
         top_p=0.95,
     )
     
-    return outputs[0]['generated_text'][len(formatted_prompt):]
+    return tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+
+# Custom CSS
+css = """
+body {
+    background: linear-gradient(135deg, #f5f7fa, #c3cfe2);
+    font-family: Arial, sans-serif;
+}
+#custom-header {
+    text-align: center;
+    background: rgba(255, 255, 255, 0.8);
+    padding: 20px;
+    border-radius: 10px;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    position: relative;
+    max-width: 800px;
+    margin: 20px auto;
+}
+#custom-header h1 {
+    color: #4A90E2;
+    font-size: 2em;
+    margin-bottom: 10px;
+}
+.llama-image {
+    position: relative;
+    transition: transform 0.3s;
+    display: inline-block;
+    margin-top: 20px;
+}
+.llama-image:hover {
+    transform: scale(1.05);
+}
+.llama-image img {
+    width: 200px;
+    border-radius: 10px;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+}
+.llama-description {
+    position: absolute;
+    bottom: -30px;
+    left: 50%;
+    transform: translateX(-50%);
+    background-color: #4A90E2;
+    color: white;
+    padding: 5px 10px;
+    border-radius: 5px;
+    opacity: 0;
+    transition: opacity 0.3s;
+    white-space: nowrap;
+}
+.llama-image:hover .llama-description {
+    opacity: 1;
+}
+.artifact {
+    position: absolute;
+    background: rgba(74, 144, 226, 0.1);
+    border-radius: 50%;
+}
+.artifact.large {
+    width: 300px;
+    height: 300px;
+    top: -50px;
+    left: -150px;
+}
+.artifact.medium {
+    width: 200px;
+    height: 200px;
+    bottom: -50px;
+    right: -100px;
+}
+.artifact.small {
+    width: 100px;
+    height: 100px;
+    top: 50%;
+    left: 50%;
+    transform: translate(-50%, -50%);
+}
+"""
 
-with gr.Blocks() as demo:
-    gr.HTML(HTML_CONTENT)
-    with gr.Row():
-        with gr.Column(scale=2):
-            prompt = gr.Textbox(label="Prompt", lines=5)
-            max_length = gr.Slider(minimum=1, maximum=500, value=128, step=1, label="Max Length")
-            temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
-            submit_button = gr.Button("Generate")
-        with gr.Column(scale=2):
-            output = gr.Textbox(label="Generated Text", lines=10)
+# Custom HTML
+custom_html = """
+<div id="custom-header">
+    <div class="artifact large"></div>
+    <div class="artifact medium"></div>
+    <div class="artifact small"></div>
     
-    submit_button.click(generate_text, inputs=[prompt, max_length, temperature], outputs=[output])
+    <h1>Llama-3.1-Storm-8B Text Generation</h1>
+    <p>Generate text using the powerful Llama-3.1-Storm-8B model. Enter a prompt and let the AI create!</p>
+    
+    <div class="llama-image">
+        <img src="https://cdn-uploads.huggingface.co/production/uploads/64c75c1237333ccfef30a602/tmOlbERGKP7JSODa6T06J.jpeg" alt="Llama">
+        <div class="llama-description">Llama-3.1-Storm-8B Model</div>
+    </div>
+</div>
+"""
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(lines=5, label="Prompt"),
+        gr.Slider(minimum=1, maximum=500, value=128, step=1, label="Max Length"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
+    ],
+    outputs=gr.Textbox(lines=10, label="Generated Text"),
+    title="Llama-3.1-Storm-8B Text Generation",
+    description="Enter a prompt to generate text using the Llama-3.1-Storm-8B model.",
+    css=css,
+    article=custom_html
+)
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the app
+iface.launch()