Update app.py

app.py

@@ -1,20 +1,37 @@
-import spaces
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
+import spaces
 import torch
-from transformers import pipeline
 
 torch.set_default_device("cuda")
 
-pipe = pipeline("text-generation", model="cognitivecomputations/dolphin-2.9.1-mixtral-1x22b")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+model_id = "cognitivecomputations/dolphin-2.9.1-mixtral-1x22b"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)
 
 @spaces.GPU(duration=120)
-def predict(message, history):
-    conv = [{"role": "system", "content": "You are Dolphin, a helpful AI assistant."}]
+def predict(input_text, history):
+    chat = []
     for item in history:
-        conv.append({"role": "user", "content": item[0]})
-        conv.append({"role": "assistant", "content": item[1]})
-    conv.append({"role": "user", "content": message})
-    generated_text = pipe(conv, max_new_tokens=1024)[0]['generated_text'][-1]['content']
-    return generated_text
+        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": input_text})
+    conv = tokenizer.apply_chat_template(chat, tokenize=False)
+    
+    inputs = tokenizer(conv, return_tensors="pt").to("cuda")
+    outputs = model.generate(**inputs, max_new_tokens=512)
+    
+    generated_text = tokenizer.batch_decode(outputs)[0]
+    return generated_text#.split("<|assistant|>")[-1]
 
-gr.ChatInterface(predict).launch()
+gr.ChatInterface(predict, theme="soft").launch()