fix docker

model.py

@@ -1,6 +1,7 @@
 import os
 import torch
-from vllm import LLM, SamplingParams
+# from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import logging
 
 # Configure logging
@@ -9,41 +10,51 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+
 class ChallengePromptGenerator:
     def __init__(
         self,
-        model_local_dir="./checkpoint-15000",
+        model_local_dir="checkpoint-15000",
     ):
-        self.generator = LLM(
-            model_local_dir, 
-        )
-
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.generator = AutoModelForCausalLM.from_pretrained(model_local_dir, device_map=self.device)
+        self.generator.to_bettertransformer()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_local_dir)
         
     def infer_prompt(
-        self,
+        self, 
         prompts,
         max_generation_length=77,
         beam_size=1,
         sampling_temperature=0.9,
         sampling_topk=1,
-        sampling_topp=1,
+        sampling_topp=1
     ):
-        added_prompts = [f"{self.generator.get_tokenizer().bos_token} {prompt}" for prompt in prompts]
+        # Add bos
+        prompts = [f"{self.tokenizer.bos_token} {prompt}" for prompt in prompts]
+        
+        # Prepare inputs
+        inputs = self.tokenizer(
+            prompts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=256,
+            add_special_tokens=False
+        ).to(self.device)
         
-        sampling_params = SamplingParams(
-            max_tokens=max_generation_length,
+        # Generate
+        outputs = self.generator.generate(
+            **inputs,
+            max_length=max_generation_length,
+            num_beams=beam_size,
             temperature=sampling_temperature,
             top_k=sampling_topk,
             top_p=sampling_topp,
-            use_beam_search=(beam_size > 1),
+            do_sample=True,
+            pad_token_id=self.tokenizer.pad_token_id
         )
         
-        outputs = self.generator.generate(added_prompts, sampling_params)
-        out = []
-        for i in range(len(outputs)):
-            tmp_out = prompts[i] + outputs[i].outputs[0].text
-            # droop last unfished sentence
-            if tmp_out[-1] != ".":
-                tmp_out = ".".join(tmp_out.split(".")[:-1])
-            out.append(tmp_out)
-        return out
+        # Decode
+        decoded_outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        return decoded_outputs

packages.txt

@@ -1,3 +0,0 @@
-gcc-12 
-g++-12 
-libnuma-dev

requirements.txt

@@ -1,7 +1,5 @@
-wheel 
-packaging 
-ninja 
-numpy
-gradio
-torch -f https://download.pytorch.org/whl/cpu
-vllm
+fastapi
+uvicorn
+pydantic
+'transformers[torch]'
+optimum