Update app.py

app.py

@@ -1,111 +1,44 @@
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# HTML template for custom UI
+HTML_TEMPLATE = """
+<style>
+    .llama-image {
+        display: flex;
+        justify-content: center;
+        margin-bottom: 20px;
+    }
+    .llama-image img {
+        max-width: 300px;
+        border-radius: 10px;
+        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
+    }
+    .llama-description {
+        text-align: center;
+        font-weight: bold;
+        margin-top: 10px;
+    }
+</style>
+<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>
+<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>
+"""
 
 # 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 template
-HTML_TEMPLATE = """
-<style>
-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 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>
-"""
-
 @spaces.GPU(duration=120)
 def generate_text(prompt, max_length, temperature):
     messages = [
@@ -114,8 +47,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,
@@ -123,8 +58,9 @@ 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)
 
+# Create Gradio interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=[
@@ -135,7 +71,12 @@ iface = gr.Interface(
     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.",
-    article=HTML_TEMPLATE
+    article=None,
+    css=".gradio-container {max-width: 800px; margin: auto;}",
 )
 
-iface.launch()
+iface.launch(
+    additional_inputs=[
+        gr.HTML(HTML_TEMPLATE)
+    ]
+)