Update app.py

app.py

@@ -32,49 +32,45 @@ system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy "
                  "while always answering questions in full first principles analysis type of thinking "
                  "without using any analogies and always showing full working code or output in his answers.")
 
-# Define a chat template
-chat_template = {
-    "system": "<|system|>{content}<|end|>",
-    "user": "<|user|>{content}<|end|>",
-    "assistant": "<|assistant|>{content}<|end|>",
-}
+# Define a chat template as a string
+chat_template = "<|system|>{system_message}<|end|><|user|>{user_message}<|end|><|assistant|>"
 
 @spaces.GPU
 def generate_response(message, history, temperature, max_new_tokens):
     if model is None or tokenizer is None:
         return "Model or tokenizer not loaded properly. Please check the logs."
     
-    messages = [{"role": "system", "content": system_prompt},
-                {"role": "user", "content": message}]
+    # Construct the full prompt
+    full_prompt = chat_template.format(system_message=system_prompt, user_message=message)
     
-    inputs = tokenizer.apply_chat_template(messages, chat_template=chat_template, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(DEVICE)
     
     with torch.no_grad():
         generate_ids = model.generate(
-            inputs,
+            inputs.input_ids,
             max_new_tokens=max_new_tokens,
             do_sample=True,
             temperature=temperature,
             eos_token_id=tokenizer.eos_token_id,
         )
-    response = tokenizer.decode(generate_ids[0, inputs.shape[1]:], skip_special_tokens=True)
+    response = tokenizer.decode(generate_ids[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return response.strip()
 
 css = """
   #output {
-    height: 900px; 
+    height: 1000px; 
     overflow: auto; 
-    border: 1px solid #ccc; 
+    border: 2px solid #ccc; 
   }
 """
 
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# Nisten's Karpathy Chatbot with OSS OLMoE")
     chatbot = gr.Chatbot(elem_id="output")
-    msg = gr.Textbox(label="Your message")
+    msg = gr.Textbox(label="Meow")
     with gr.Row():
         temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
-        max_new_tokens = gr.Slider(minimum=50, maximum=4000, value=1000, step=50, label="Max New Tokens")
+        max_new_tokens = gr.Slider(minimum=50, maximum=4000, value=2000, step=50, label="Max New Tokens")
     clear = gr.Button("Clear")
 
     def user(user_message, history):
@@ -93,4 +89,4 @@ with gr.Blocks(css=css) as demo:
 
 if __name__ == "__main__":
     demo.queue(api_open=True)
-    demo.launch(debug=True, show_api=True, share=True)
+    demo.launch(debug=True, show_api=True, share=True)