Update app.py

app.py

@@ -1,6 +1,11 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+    BitsAndBytesConfig,
+)
 import os
 from threading import Thread
 import spaces
@@ -9,27 +14,27 @@ import time
 token = os.environ["HF_TOKEN"]
 
 quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16
+    load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16
 )
 
-model = AutoModelForCausalLM.from_pretrained("google/gemma-1.1-7b-it", 
-                                             quantization_config=quantization_config,
-                                             token=token)
+model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-1.1-7b-it", quantization_config=quantization_config, token=token
+)
 tok = AutoTokenizer.from_pretrained("google/gemma-1.1-7b-it", token=token)
 
 if torch.cuda.is_available():
-    device = torch.device('cuda')
+    device = torch.device("cuda")
     print(f"Using GPU: {torch.cuda.get_device_name(device)}")
 else:
-    device = torch.device('cpu')
+    device = torch.device("cpu")
     print("Using CPU")
-    
+
 # model = model.to(device)
-    # Dispatch Errors
+# Dispatch Errors
+
 
 @spaces.GPU
-def chat(message, history):
+def chat(message, history, temperature, top_p, top_k, max_tokens):
     start_time = time.time()
     chat = []
     for item in history:
@@ -40,16 +45,16 @@ def chat(message, history):
     messages = tok.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
     model_inputs = tok([messages], return_tensors="pt").to(device)
     streamer = TextIteratorStreamer(
-        tok, timeout=10., skip_prompt=True, skip_special_tokens=True)
+        tok, timeout=10.0, skip_prompt=True, skip_special_tokens=True
+    )
     generate_kwargs = dict(
         model_inputs,
         streamer=streamer,
-        max_new_tokens=1024,
+        max_new_tokens=max_tokens,
         do_sample=True,
-        top_p=0.95,
-        top_k=1000,
-        temperature=0.75,
-        num_beams=1,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
@@ -66,9 +71,36 @@ def chat(message, history):
     tokens = len(tok.tokenize(partial_text))
     tokens_per_second = tokens / total_time if total_time > 0 else 0
 
-    # Append the timing information to the final output
-    timing_info = f"\nTime taken to first token: {first_token_time:.2f} seconds\nTokens per second: {tokens_per_second:.2f}"
+    timing_info = f"\n\nTime taken to first token: {first_token_time:.2f} seconds\nTokens per second: {tokens_per_second:.2f}"
     yield partial_text + timing_info
 
-demo = gr.ChatInterface(fn=chat, examples=[["Write me a poem about Machine Learning."]], title="Chat With LLMS")
+
+demo = gr.ChatInterface(
+    fn=chat,
+    examples=[["Write me a poem about Machine Learning."]],
+    additional_inputs_accordion=gr.Accordion(
+        label="⚙️ Parameters", open=False, render=False
+    ),
+    additional_inputs=[
+        gr.Slider(
+            minimum=0, maximum=1, step=0.1, value=0.9, label="Temperature", render=False
+        ),
+        gr.Slider(
+            minimum=0, maximum=1, step=0.1, value=0.95, label="top_p", render=False
+        ),
+        gr.Slider(
+            minimum=1, maximum=10000, step=5, value=1000, label="top_k", render=False
+        ),
+        gr.Slider(
+            minimum=128,
+            maximum=4096,
+            step=1,
+            value=1024,
+            label="Max new tokens",
+            render=False,
+        ),
+    ],
+    multimodal=False,
+    title="Chat With LLMs",
+)
 demo.launch()