Update model.py

model.py

@@ -1,9 +1,11 @@
 from threading import Thread
 from typing import Iterator
 
-
+import torch
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-model="/kaggle/input/deberta-v3-large-hf-weights"
+
+model_id = 'meta-llama/Llama-2-13b-chat-hf'
+
 if torch.cuda.is_available():
     config = AutoConfig.from_pretrained(model_id)
     config.pretraining_tp = 1
@@ -17,19 +19,28 @@ if torch.cuda.is_available():
 else:
     model = None
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-def preprocess(sample):
-    first_sentences = [sample["prompt"]] * 5
-    second_sentences = [sample[option] for option in "ABCDE"]
-    tokenized_sentences = tokenizer(first_sentences, second_sentences, truncation=True, padding=True, return_tensors="pt")
-    sample["input_ids"] = tokenized_sentences["input_ids"]
-    sample["attention_mask"] = tokenized_sentences["attention_mask"]
-    return sample
+
+
+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']
+    # 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] ')
+    message = message.strip() if do_strip else message
+    texts.append(f'{message} [/INST]')
+    return ''.join(texts)
+
 
 def get_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> int:
     prompt = get_prompt(message, chat_history, system_prompt)
     input_ids = tokenizer([prompt], return_tensors='np', add_special_tokens=False)['input_ids']
     return input_ids.shape[-1]
 
+
 def run(message: str,
         chat_history: list[tuple[str, str]],
         system_prompt: str,
@@ -61,3 +72,4 @@ def run(message: str,
     for text in streamer:
         outputs.append(text)
         yield ''.join(outputs)
+