Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import OlmoeForCausalLM, AutoTokenizer
 import torch
 import subprocess
 import sys
+import os
 
 # Force upgrade transformers to the latest version
 subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "transformers"])
@@ -36,25 +37,23 @@ 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}]
-    
+    messages = [{"role": "user", "content": message}]
     inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
     
     with torch.no_grad():
         generate_ids = model.generate(
-            inputs,
+            **inputs,
             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], skip_special_tokens=True)
     return response.strip()
 
 css = """
   #output {
-    height: 900px; 
+    height: 500px; 
     overflow: auto; 
     border: 1px solid #ccc; 
   }
@@ -85,4 +84,4 @@ with gr.Blocks(css=css) as demo:
 
 if __name__ == "__main__":
     demo.queue(api_open=False)
-    demo.launch(debug=True, show_api=True, share=True)
+    demo.launch(debug=True, show_api=False, share=True )