Update model.py

model.py

@@ -1,25 +1,19 @@
-from threading import Thread
+import os
 from typing import Iterator
 
-import torch
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from text_generation import Client
 
-model_id = 'codellama/CodeLlama-13b-Instruct-hf'
+model_id = 'codellama/CodeLlama-34b-Instruct-hf'
 
-if torch.cuda.is_available():
-    config = AutoConfig.from_pretrained(model_id)
-    config.pretraining_tp = 1
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        config=config,
-        torch_dtype=torch.float16,
-        load_in_4bit=True,
-        device_map='auto',
-        use_safetensors=False,
-    )
-else:
-    model = None
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+API_URL = "https://api-inference.huggingface.co/models/" + model_id
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+client = Client(
+    API_URL,
+    headers={"Authorization": f"Bearer {HF_TOKEN}"},
+)
+EOS_STRING = "</s>"
+EOT_STRING = "<EOT>"
 
 
 def get_prompt(message: str, chat_history: list[tuple[str, str]],
@@ -36,12 +30,6 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
     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,
@@ -50,26 +38,20 @@ def run(message: str,
         top_p: float = 0.9,
         top_k: int = 50) -> Iterator[str]:
     prompt = get_prompt(message, chat_history, system_prompt)
-    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')
 
-    streamer = TextIteratorStreamer(tokenizer,
-                                    timeout=10.,
-                                    skip_prompt=True,
-                                    skip_special_tokens=True)
     generate_kwargs = dict(
-        inputs,
-        streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
-        num_beams=1,
     )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield ''.join(outputs)
+    stream = client.generate_stream(prompt, **generate_kwargs)
+    output = ""
+    for response in stream:
+        if any([end_token in response.token.text for end_token in [EOS_STRING, EOT_STRING]]):
+            return output
+        else:
+            output += response.token.text
+        yield output
+    return output