Update app.py

app.py

@@ -8,7 +8,7 @@ import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-MAX_MAX_NEW_TOKENS = 2048
+MAX_MAX_NEW_TOKENS = 128
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 model = None
@@ -19,7 +19,8 @@ def load_model():
     model_id = "stabilityai/ar-stablelm-2-chat"
     model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
     tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    tokenizer.use_default_system_prompt = False
+    model.generation_config.pad_token_id = model.generation_config.eos_token_id
+
 
 
 def generate(
@@ -49,8 +50,11 @@ def generate(
         {"input_ids": input_ids},
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        do_sample=False,
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,  # Stop generation at <EOS>
         temperature=temperature,
+        top_p=top_p,
+        top_k=top_k
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
@@ -115,7 +119,7 @@ with gr.Blocks(css_paths="style.css", fill_height=True) as demo:
         try:
             login(token = token)
             load_model()
-            return f"Authenticated successfully"
+            return "Authenticated successfully"
         except:
             return "Invalid token. Please try again."
 
@@ -129,4 +133,4 @@ with gr.Blocks(css_paths="style.css", fill_height=True) as demo:
     
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()