Update app.py

app.py

@@ -1,88 +1,111 @@
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, pipeline
 
-# HTML template for custom UI
+# 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,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+# HTML template
 HTML_TEMPLATE = """
 <style>
-    body { background: linear-gradient(135deg, #f5f7fa, #c3cfe2); }
-    #app-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;
-    }
-    #app-header h1 { color: #4CAF50; font-size: 2em; margin-bottom: 10px; }
-    .concept { position: relative; transition: transform 0.3s; }
-    .concept:hover { transform: scale(1.1); }
-    .concept img {
-        width: 100px;
-        border-radius: 10px;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-    }
-    .concept-description {
-        position: absolute;
-        bottom: -30px;
-        left: 50%;
-        transform: translateX(-50%);
-        background-color: #4CAF50;
-        color: white;
-        padding: 5px 10px;
-        border-radius: 5px;
-        opacity: 0;
-        transition: opacity 0.3s;
-    }
-    .concept:hover .concept-description { opacity: 1; }
-    .artifact {
-        position: absolute;
-        background: rgba(76, 175, 80, 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%);
-    }
+body {
+    background: linear-gradient(135deg, #f5f7fa, #c3cfe2);
+    font-family: Arial, sans-serif;
+}
+#app-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;
+}
+#app-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%);
+}
 </style>
+
 <div id="app-header">
     <div class="artifact large"></div>
     <div class="artifact medium"></div>
     <div class="artifact small"></div>
-    <h1>Llama-3.1-Storm-8B Text Generator</h1>
-    <p>Generate text using the Llama-3.1-Storm-8B model by providing a prompt.</p>
-    <div style="display: flex; justify-content: center; gap: 20px; margin-top: 20px;">
-        <div class="concept">
-            <img src="https://raw.githubusercontent.com/huggingface/huggingface.js/main/packages/inference/src/tasks/images/llama.png" alt="Llama">
-            <div class="concept-description">Llama Model</div>
-        </div>
-        <div class="concept">
-            <img src="https://raw.githubusercontent.com/huggingface/huggingface.js/main/packages/inference/src/tasks/images/language.png" alt="Language">
-            <div class="concept-description">Natural Language Processing</div>
-        </div>
-        <div class="concept">
-            <img src="https://raw.githubusercontent.com/huggingface/huggingface.js/main/packages/inference/src/tasks/images/text-generation.png" alt="Text Generation">
-            <div class="concept-description">Text Generation</div>
-        </div>
+    
+    <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>
 """
 
-# Load the model and tokenizer
-model_name = "akjindal53244/Llama-3.1-Storm-8B"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
-)
-
 @spaces.GPU(duration=120)
 def generate_text(prompt, max_length, temperature):
     messages = [
@@ -91,10 +114,8 @@ def generate_text(prompt, max_length, temperature):
     ]
     formatted_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
     
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
-    
-    outputs = model.generate(
-        **inputs,
+    outputs = pipe(
+        formatted_prompt,
         max_new_tokens=max_length,
         do_sample=True,
         temperature=temperature,
@@ -102,9 +123,8 @@ def generate_text(prompt, max_length, temperature):
         top_p=0.95,
     )
     
-    return tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    return outputs[0]['generated_text'][len(formatted_prompt):]
 
-# Create Gradio interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=[