add demo

app.py

@@ -3,6 +3,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 from datetime import datetime
 
+# Model initialization
 model_id = "BSC-LT/salamandra-2b-instruct"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -10,9 +11,13 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.bfloat16
 )
 
+# Set pad_token_id to eos_token_id if it's not set
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
 description = """
 Salamandra-2b-instruct is a Transformer-based decoder-only language model that has been pre-trained on 7.8 trillion tokens of highly curated data. 
 The pre-training corpus contains text in 35 European languages and code. This instruction-tuned variant can be used as a general-purpose assistant.
@@ -27,36 +32,42 @@ On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [Buil
 🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
 """
 
-def generate_text(prompt, temperature, max_new_tokens, top_p, repetition_penalty):
+def generate_text(system_prompt, user_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
     date_string = datetime.today().strftime('%Y-%m-%d')
-    message = [{"role": "user", "content": prompt}]
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt}
+    ]
     
     chat_prompt = tokenizer.apply_chat_template(
-        message,
+        messages,
         tokenize=False,
         add_generation_prompt=True,
         date_string=date_string
     )
     
-    inputs = tokenizer.encode(chat_prompt, add_special_tokens=False, return_tensors="pt")
+    inputs = tokenizer(chat_prompt, return_tensors="pt", padding=True, truncation=True)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
     
     outputs = model.generate(
-        input_ids=inputs.to(model.device),
+        **inputs,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         top_p=top_p,
         repetition_penalty=repetition_penalty,
-        do_sample=True
+        do_sample=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
     )
     
     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return generated_text.split("assistant\n")[-1].strip()
 
-def update_output(prompt, temperature, max_new_tokens, top_p, repetition_penalty):
-    return generate_text(prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+def update_output(system_prompt, user_prompt, temperature, max_new_tokens, top_p, repetition_penalty):
+    return generate_text(system_prompt, user_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🙋🏻‍♂️ Welcome to Tonic's 📲🦎Salamandra-2b-instruct Demo")
+    gr.Markdown("# 🦎 Welcome to Tonic's Salamandra-2b-instruct Demo")
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -66,8 +77,13 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column(scale=1):
-            prompt = gr.Textbox(lines=5, label="🙋‍♂️ Input Prompt")
-            generate_button = gr.Button("Try 📲🦎Salamandra-2b-instruct")
+            system_prompt = gr.Textbox(
+                lines=3, 
+                label="🖥️ System Prompt", 
+                value="You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information."
+            )
+            user_prompt = gr.Textbox(lines=5, label="🙋‍♂️ User Prompt")
+            generate_button = gr.Button("Generate with 🦎 Salamandra-2b-instruct")
             
             with gr.Accordion("🧪 Parameters", open=False):
                 temperature = gr.Slider(0.0, 1.0, value=0.7, label="🌡️ Temperature")
@@ -76,24 +92,24 @@ with gr.Blocks() as demo:
                 repetition_penalty = gr.Slider(1.0, 2.0, value=1.2, label="🔁 Repetition Penalty")
         
         with gr.Column(scale=1):
-            output = gr.Textbox(lines=10, label="📲🦎Salamandra")
+            output = gr.Textbox(lines=10, label="🦎 Salamandra-2b-instruct Output")
     
     generate_button.click(
         update_output,
-        inputs=[prompt, temperature, max_new_tokens, top_p, repetition_penalty],
+        inputs=[system_prompt, user_prompt, temperature, max_new_tokens, top_p, repetition_penalty],
         outputs=output
     )
     
     gr.Examples(
         examples=[
-            ["What are the main advantages of living in a big city like Barcelona?"],
-            ["Explain the process of photosynthesis in simple terms."],
-            ["What are some effective strategies for learning a new language?"],
-            ["Describe the potential impacts of artificial intelligence on the job market in the next decade."],
-            ["What are the key differences between renewable and non-renewable energy sources?"]
+            ["You are a helpful assistant.", "What are the main advantages of living in a big city like Barcelona?"],
+            ["You are a biology teacher explaining concepts to students.", "Explain the process of photosynthesis in simple terms."],
+            ["You are a language learning expert.", "What are some effective strategies for learning a new language?"],
+            ["You are an AI and technology expert.", "Describe the potential impacts of artificial intelligence on the job market in the next decade."],
+            ["You are an environmental scientist.", "What are the key differences between renewable and non-renewable energy sources?"]
         ],
-        inputs=prompt,
-        outputs=prompt,
+        inputs=[system_prompt, user_prompt],
+        outputs=[system_prompt, user_prompt],
         label="Example Prompts"
     )