Update app.py (#1)

- Update app.py (a4a84f65e02a927e0d3994287a1a8706c15d63a4)

app.py

@@ -1,94 +1,75 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+    BitsAndBytesConfig,
+)
 import os
 from threading import Thread
 import spaces
 import time
-import subprocess
 
-MIN_TOKENS=128
-MAX_TOKENS=8192
-DEFAULT_TOKENS=2048
-DURATION=60
+token = os.environ["HF_TOKEN"]
 
-# Install flash attention
-subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load model and tokenizer once when the app starts
-model_token = os.environ["HF_TOKEN"]
 model = AutoModelForCausalLM.from_pretrained(
-    "microsoft/Phi-3-mini-128k-instruct",
-    token=model_token,
-    trust_remote_code=True,
+    "microsoft/Phi-3-mini-128k-instruct", quantization_config=quantization_config, token=token
 )
-tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct", token=model_token)
-
-# Set device (GPU or CPU)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-# Define error handling function
-def handle_error(error):
-    return {"error": str(error)}
-
-# Define chat function with input validation and error handling
-@spaces.GPU(duration=DURATION)
-def chat(message, history, temperature, do_sample, max_tokens):
-    try:
-        # Validate input
-        if not message:
-            raise ValueError("Please enter a message")
-        if temperature < 0 or temperature > 1:
-            raise ValueError("Temperature must be between 0 and 1")
-        if max_tokens < MIN_TOKENS or max_tokens > MAX_TOKENS:
-            raise ValueError(f"Max tokens must be between {MIN_TOKENS} and {MAX_TOKENS}")
+tok = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct", token=token)
 
-        # Prepare chat history
-        chat = []
-        for item in history:
-            chat.append({"role": "user", "content": item[0]})
-            if item[1] is not None:
-                chat.append({"role": "assistant", "content": item[1]})
-        chat.append({"role": "user", "content": message})
 
-        # Generate response
-        messages = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-        model_inputs = tokenizer([messages], return_tensors="pt").to(device)
-        streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-        generate_kwargs = dict(
-            model_inputs,
-            streamer=streamer,
-            max_new_tokens=max_tokens,
-            do_sample=do_sample,
-            temperature=temperature,
-            eos_token_id=[tokenizer.eos_token_id],
-        )
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print(f"Using GPU: {torch.cuda.get_device_name(device)}")
+else:
+    device = torch.device("cpu")
+    print("Using CPU")
 
-        # Generate response
-        t = Thread(target=model.generate, kwargs=generate_kwargs)
-        t.start()
+@spaces.GPU(duration=150)
+def chat(message, history, temperature,do_sample, max_tokens):
+    chat = []
+    for item in history:
+        chat.append({"role": "user", "content": item[0]})
+        if item[1] is not None:
+            chat.append({"role": "assistant", "content": item[1]})
+    chat.append({"role": "user", "content": message})
+    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.0, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+    )
+    
+    if temperature == 0:
+        generate_kwargs['do_sample'] = False
+    
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
 
-        # Yield partial responses
-        partial_text = ""
-        for new_text in streamer:
-            partial_text += new_text
-            yield partial_text
-
-        # Yield final response
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
         yield partial_text
 
-    except Exception as e:
-        yield handle_error(e)
+    tokens = len(tok.tokenize(partial_text))
+    yield partial_text 
+
 
-# Create Gradio interface
 demo = gr.ChatInterface(
     fn=chat,
     examples=[["Write me a poem about Machine Learning."]],
+    # multimodal=False,
     additional_inputs_accordion=gr.Accordion(
         label="⚙️ Parameters", open=False, render=False
     ),
@@ -96,20 +77,18 @@ demo = gr.ChatInterface(
         gr.Slider(
             minimum=0, maximum=1, step=0.1, value=0.9, label="Temperature", render=False
         ),
-        gr.Checkbox(label="Sampling", value=True),
+        gr.Checkbox(label="Sampling",value=True),
         gr.Slider(
-            minimum=MIN_TOKENS,
-            maximum=MAX_TOKENS,
+            minimum=128,
+            maximum=4096,
             step=1,
-            value=DEFAULT_TOKENS,
+            value=512,
             label="Max new tokens",
             render=False,
         ),
     ],
     stop_btn="Stop Generation",
     title="Chat With LLMs",
-    description="Now Running [microsoft/Phi-3-mini-128k-instruct](https://huggingface.co/microsoft/Phi-3-mini-128k-instruct)",
+    description="Now Running [Microsoft Phi-3-mini-128k-instruct](https://huggingface.co/microsoft/Phi-3-mini-128k-instruct) in 4bit"
 )
-
-# Launch Gradio app
 demo.launch()