Update model.py

model.py

@@ -3,31 +3,42 @@ from typing import Iterator
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 
-model_id = 'meta-llama/Llama-2-7b-chat-hf'
+model_id = "TheBloke/Chronos-Beluga-v2-13B-GPTQ"
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
+model = AutoGPTQForCausalLM.from_quantized(model_id,
+        use_safetensors=True,
+        trust_remote_code=False,
+        device="cuda:0",
+        use_triton=False,
+        quantize_config=None)
 
-if torch.cuda.is_available():
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        device_map='auto'
-    )
-else:
-    model = None
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+# model_id = 'meta-llama/Llama-2-7b-chat-hf'
+
+# if torch.cuda.is_available():
+#     model = AutoModelForCausalLM.from_pretrained(
+#         model_id,
+#         torch_dtype=torch.float16,
+#         device_map='auto'
+#     )
+# else:
+#     model = None
+# tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
 def get_prompt(message: str, chat_history: list[tuple[str, str]],
                system_prompt: str) -> str:
-    texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']
+    # texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']
+    texts = [f'{system_prompt}\n\n']
     # The first user input is _not_ stripped
     do_strip = False
     for user_input, response in chat_history:
         user_input = user_input.strip() if do_strip else user_input
         do_strip = True
-        texts.append(f'{user_input} [/INST] {response.strip()} </s><s>[INST] ')
+        texts.append(f'{user_input} {response.strip()} ')
     message = message.strip() if do_strip else message
-    texts.append(f'{message} [/INST]')
+    texts.append(f'{message}')
     return ''.join(texts)