add vllm

README.md

@@ -78,4 +78,4 @@ For bug fixes or questions, either open an issue or create a branch prefixed wit
 
 ## Accknowledgement
 
-Thanks for the GPU grant from Huggingface.
+Thanks for the GPU grant from Huggingface.

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 
 from pages.arena import create_arena
-from pages.summarization_playground import create_summarization_interface
-from pages.leaderboard import create_leaderboard
 from pages.batch_evaluation import create_batch_evaluation_interface
+from pages.leaderboard import create_leaderboard
+from pages.summarization_playground import create_summarization_interface
+
 
 def welcome_message():
     return """## Clinical Dialogue Summarization

pages/arena.py

@@ -1,11 +1,12 @@
+import json
 import random
+
 import gradio as gr
-import json
 
+from pages.summarization_playground import custom_css, get_model_batch_generation
 from utils.data import dataset
 from utils.multiple_stream import stream_data
-from pages.summarization_playground import get_model_batch_generation
-from pages.summarization_playground import custom_css
+
 
 def random_data_selection():
     datapoint = random.choice(dataset)

pages/batch_evaluation.py

@@ -1,17 +1,14 @@
-from dotenv import load_dotenv
-import gradio as gr
-
-import json
 import html
+import json
 import logging
 
+import gradio as gr
 import numpy as np
+from dotenv import load_dotenv
 
-from utils.model import Model
+from pages.summarization_playground import custom_css, generate_answer
 from utils.metric import metric_rouge_score
-
-from pages.summarization_playground import generate_answer
-from pages.summarization_playground import custom_css
+from utils.model import Model
 
 load_dotenv()
 

pages/leaderboard.py

@@ -1,7 +1,9 @@
 import html
 import json
-import pandas as pd
+
 import gradio as gr
+import pandas as pd
+
 
 # Function to create HTML tooltips
 def create_html_with_tooltip(id, base_url):

pages/summarization_playground.py

@@ -1,14 +1,13 @@
-from dotenv import load_dotenv
-import gradio as gr
+import gc
+import logging
 import random
 
-from utils.model import Model
-from utils.data import dataset
-
-import gc
+import gradio as gr
 import torch
+from dotenv import load_dotenv
 
-import logging
+from utils.data import dataset
+from utils.model import Model
 
 load_dotenv()
 

prompt/prompt.ipynb

@@ -15,7 +15,7 @@
     "        \"author\": \"Shunxi Wu\",\n",
     "        \"metric\": {\n",
     "            \"Rouge\": 0.14,\n",
-    "            \"winning_number\": 10\n",
+    "            \"winning_number\": 11\n",
     "        },\n",
     "        \"url\": \"https://docs.google.com/spreadsheets/d/1ui9ccRkzeMWAiJiRgr2ClpYTAK4uFhX44aXi0WDJY8Q/edit?gid=1699794338#gid=1699794338&range=D2\"\n",
     "    },\n",

requirements.txt

@@ -9,4 +9,5 @@ torchvision
 torchaudio
 datasets
 rouge-score
-markdown
+markdown
+vllm

utils/model.py

@@ -1,12 +1,12 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
-import torch
+import logging
+import os
 
+import torch
 from huggingface_hub import login
-import os 
+from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from vllm import LLM, SamplingParams
 
-import logging
-
-login(token = os.getenv('HF_TOKEN'))
+login(token=os.getenv('HF_TOKEN'))
 
 class Model(torch.nn.Module):
     number_of_models = 0
@@ -23,89 +23,106 @@ class Model(torch.nn.Module):
         
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.name = model_name
+        self.use_vllm = model_name != "google-t5/t5-large"
 
-        logging.info(f'start loading model {self.name}')
+        logging.info(f'Start loading model {self.name}')
 
-        if model_name == "google-t5/t5-large":
-            # For T5 or any other Seq2Seq model
-            self.model = AutoModelForSeq2SeqLM.from_pretrained(
-                model_name, torch_dtype=torch.bfloat16, device_map="auto"
+        if self.use_vllm:
+            # 使用vLLM加载模型
+            self.llm = LLM(
+                model=model_name,
+                dtype="bfloat16",
+                tokenizer=model_name,
+                trust_remote_code=True
             )
         else:
-            # For GPT-like models or other causal language models
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name, torch_dtype=torch.bfloat16, device_map="auto"
+            # 加载原始transformers模型
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(
+                model_name, 
+                torch_dtype=torch.bfloat16, 
+                device_map="auto"
             )
+            self.model.eval()
 
         logging.info(f'Loaded model {self.name}')
-
-        self.model.eval()
         self.update()
 
     @classmethod
     def update(cls):
         cls.number_of_models += 1
 
-    def return_mode_name(self):
-        return self.name
-    
-    def return_tokenizer(self):
-        return self.tokenizer
-    
-    def return_model(self):
-        return self.model
-    
-    def streaming(self, content_list, temp=0.001, max_length=500, do_sample=True):
-        # Convert list of texts to input IDs
-        input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
-
-        # Set up the initial generation parameters
-        gen_kwargs = {
-            "input_ids": input_ids,
-            "do_sample": do_sample,
-            "temperature": temp,
-            "eos_token_id": self.tokenizer.eos_token_id,
-            "max_new_tokens": 1,  # Generate one token at a time
-            "return_dict_in_generate": True,
-            "output_scores": True
-        }
-
-        # Generate and yield tokens one by one
-        generated_tokens = 0
-        batch_size = input_ids.shape[0]
-        active_sequences = torch.arange(batch_size)
+    def gen(self, content_list, temp=0.001, max_length=500, do_sample=True):
+        if self.use_vllm:
+            sampling_params = SamplingParams(
+                temperature=temp,
+                max_tokens=max_length,
+                top_p=0.95 if do_sample else 1.0,
+                stop_token_ids=[self.tokenizer.eos_token_id]
+            )
+            outputs = self.llm.generate(content_list, sampling_params)
+            return [output.outputs[0].text for output in outputs]
+        else:
+            input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
+            outputs = self.model.generate(
+                input_ids,
+                max_new_tokens=max_length,
+                do_sample=do_sample,
+                temperature=temp,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )
+            return self.tokenizer.batch_decode(outputs[:, input_ids.shape[1]:], skip_special_tokens=True)
 
-        while generated_tokens < max_length and len(active_sequences) > 0:
-            with torch.no_grad():
-                output = self.model.generate(**gen_kwargs)
+    def streaming(self, content_list, temp=0.001, max_length=500, do_sample=True):
+        if self.use_vllm:
+            sampling_params = SamplingParams(
+                temperature=temp,
+                max_tokens=max_length,
+                top_p=0.95 if do_sample else 1.0,
+                stop_token_ids=[self.tokenizer.eos_token_id]
+            )
+            outputs = self.llm.generate(content_list, sampling_params, stream=True)
             
-            next_tokens = output.sequences[:, -1].unsqueeze(-1)
+            prev_token_ids = [[] for _ in content_list]
             
-            # Yield the newly generated tokens for each sequence in the batch
-            for i, token in zip(active_sequences, next_tokens):
-                yield i, self.tokenizer.decode(token[0], skip_special_tokens=True)
-
-            # Update input_ids for the next iteration
-            gen_kwargs["input_ids"] = torch.cat([gen_kwargs["input_ids"], next_tokens], dim=-1)
-            generated_tokens += 1
-
-            # Check for completed sequences
-            completed = (next_tokens.squeeze(-1) == self.tokenizer.eos_token_id).nonzero().squeeze(-1)
-            active_sequences = torch.tensor([i for i in active_sequences if i not in completed])
-            if len(active_sequences) > 0:
-                gen_kwargs["input_ids"] = gen_kwargs["input_ids"][active_sequences]
-
-
-    def gen(self, content_list, temp=0.001, max_length=500, do_sample=True):
-        # Convert list of texts to input IDs
-        input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
-
-        # Non-streaming generation (unchanged)
-        outputs = self.model.generate(
-            input_ids,
-            max_new_tokens=max_length,
-            do_sample=do_sample,
-            temperature=temp,
-            eos_token_id=self.tokenizer.eos_token_id,
-        )
-        return self.tokenizer.batch_decode(outputs[:, input_ids.shape[1]:], skip_special_tokens=True)
+            for output in outputs:
+                for i, request_output in enumerate(output.outputs):
+                    current_token_ids = request_output.token_ids
+                    new_token_ids = current_token_ids[len(prev_token_ids[i]):]
+                    prev_token_ids[i] = current_token_ids.copy()
+                    
+                    for token_id in new_token_ids:
+                        token_text = self.tokenizer.decode(token_id, skip_special_tokens=True)
+                        yield i, token_text
+        else:
+            input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
+            
+            gen_kwargs = {
+                "input_ids": input_ids,
+                "do_sample": do_sample,
+                "temperature": temp,
+                "eos_token_id": self.tokenizer.eos_token_id,
+                "max_new_tokens": 1,
+                "return_dict_in_generate": True,
+                "output_scores": True
+            }
+
+            generated_tokens = 0
+            batch_size = input_ids.shape[0]
+            active_sequences = torch.arange(batch_size)
+
+            while generated_tokens < max_length and len(active_sequences) > 0:
+                with torch.no_grad():
+                    output = self.model.generate(**gen_kwargs)
+                
+                next_tokens = output.sequences[:, -1].unsqueeze(-1)
+                
+                for i, token in zip(active_sequences, next_tokens):
+                    yield i.item(), self.tokenizer.decode(token[0], skip_special_tokens=True)
+
+                gen_kwargs["input_ids"] = torch.cat([gen_kwargs["input_ids"], next_tokens], dim=-1)
+                generated_tokens += 1
+
+                completed = (next_tokens.squeeze(-1) == self.tokenizer.eos_token_id).nonzero().squeeze(-1)
+                active_sequences = torch.tensor([i for i in active_sequences if i not in completed])
+                if len(active_sequences) > 0:
+                    gen_kwargs["input_ids"] = gen_kwargs["input_ids"][active_sequences]

utils/multiple_stream.py

@@ -1,5 +1,6 @@
 import copy
 import random
+
 import gradio as gr
 
 TEST = """ Test of Time. A Benchmark for Evaluating LLMs on Temporal Reasoning. Large language models (LLMs) have