crfm-helm 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

helm/{proxy/clients → clients}/huggingface_client.py

@@ -5,7 +5,7 @@ from transformers.generation.stopping_criteria import (
     StoppingCriteria,
     StoppingCriteriaList,
 )
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, TypedDict
 
 from helm.common.cache import CacheConfig
 from helm.common.hierarchical_logger import htrack_block, hlog
@@ -14,11 +14,11 @@ from helm.common.request import (
     EMBEDDING_UNAVAILABLE_REQUEST_RESULT,
     Request,
     RequestResult,
-    Sequence,
+    GeneratedOutput,
     Token,
 )
 from .client import CachingClient, truncate_sequence
-from helm.proxy.tokenizers.huggingface_tokenizer import HuggingFaceTokenizer, WrappedPreTrainedTokenizer, resolve_alias
+from helm.tokenizers.huggingface_tokenizer import HuggingFaceTokenizer, WrappedPreTrainedTokenizer
 from threading import Lock
 
 
@@ -36,6 +36,20 @@ class StopAtSpecificTokenCriteria(StoppingCriteria):
         return bool(torch.all(current_sequence == stop_sequence_tensor).item())
 
 
+class HuggingFaceRequest(TypedDict):
+    """Data passed between make_request and serve_request. Used as the cache key."""
+
+    engine: str
+    prompt: str
+    temperature: float
+    num_return_sequences: int
+    max_new_tokens: int
+    top_p: float
+    echo_prompt: bool
+    top_k_per_token: int
+    stop_sequences: List
+
+
 class HuggingFaceServer:
     """A thin wrapper around a Hugging Face AutoModelForCausalLM for HuggingFaceClient to call."""
 
@@ -55,30 +69,24 @@ class HuggingFaceServer:
                 pretrained_model_name_or_path, **kwargs
             )
 
-    def serve_request(self, raw_request: Dict[str, Any]):
+    def serve_request(self, raw_request: HuggingFaceRequest) -> Dict:
         with self.wrapped_tokenizer as tokenizer:
             encoded_input = tokenizer(raw_request["prompt"], return_tensors="pt", return_token_type_ids=False).to(
                 self.device
             )
-        raw_request = deepcopy(raw_request)
-        raw_request["do_sample"] = True
-        raw_request["return_dict_in_generate"] = True
-        raw_request["output_scores"] = True
-        top_k_per_token: int = raw_request["top_k_per_token"]
-        del raw_request["top_k_per_token"]
         stopping_criteria: Optional[StoppingCriteriaList] = None
+        optional_args = {}
         if len(raw_request["stop_sequences"]) > 0:
             with self.wrapped_tokenizer as tokenizer:
                 stop_sequence_ids = tokenizer(
                     raw_request["stop_sequences"], return_token_type_ids=False, add_special_tokens=False
                 )
             if len(stop_sequence_ids.input_ids) == 1 and len(stop_sequence_ids.input_ids[0]) == 1:
-                raw_request["eos_token_id"] = stop_sequence_ids.input_ids[0][0]
+                optional_args["eos_token_id"] = stop_sequence_ids.input_ids[0][0]
             else:
                 stopping_criteria = StoppingCriteriaList()
                 for stop_sequence_input_ids in stop_sequence_ids.input_ids:
                     stopping_criteria.append(StopAtSpecificTokenCriteria(stop_sequence=stop_sequence_input_ids))
-            del raw_request["stop_sequences"]
 
         # Check if we need to compute the perplexity of the prompt (#1497)
         compute_logprobs_only = (
@@ -94,64 +102,42 @@ class HuggingFaceServer:
             sequences = encoded_input["input_ids"]
             scores = output.logits
         else:
-            # Strip out irrelevant parameters
-            relevant_raw_request = {
-                key: raw_request[key]
-                for key in raw_request
-                if key not in ["engine", "prompt", "echo_prompt", "stop_sequences"]
-            }
-
             output = self.model.generate(
                 **encoded_input,
-                **relevant_raw_request,
+                temperature=raw_request["temperature"],
+                num_return_sequences=raw_request["num_return_sequences"],
+                max_new_tokens=raw_request["max_new_tokens"],
+                top_p=raw_request["top_p"],
+                do_sample=True,
+                return_dict_in_generate=True,
+                output_scores=True,
+                **optional_args,
                 stopping_criteria=stopping_criteria,
             )
             sequences = output.sequences
             scores = output.scores
 
         prompt_tokens_logprobs = []
-        prompt_tokens_top_logprobs_dicts: List[Dict] = []
         if compute_logprobs_only:
             # Append the logprob of the first token of the prompt.
             prompt_tokens_logprobs.append(0.0)
-            prompt_tokens_top_logprobs_dicts.append({})
 
             # Compute logprobs of prompt tokens.
             for completion_id in range(raw_request["num_return_sequences"]):
                 for i in range(len(sequences[completion_id]) - 1):
                     logprobs = torch.nn.functional.log_softmax(scores[completion_id][i], dim=0)
-                    topk_logprobs = torch.topk(logprobs, k=top_k_per_token)
-                    with self.wrapped_tokenizer as tokenizer:
-                        prompt_tokens_top_logprobs_dicts.append(
-                            {
-                                tokenizer.convert_ids_to_tokens(k.item()): v.item()
-                                for (k, v) in zip(topk_logprobs.indices, topk_logprobs.values)
-                            }
-                        )
                     prompt_tokens_logprobs.append(logprobs[sequences[completion_id][i + 1]].item())
 
         # Compute logprobs of generated tokens for each completed sequence.
         all_generated_tokens_logprobs = []
-        all_generated_tokens_top_logprobs_dicts = []
         for completion_id in range(raw_request["num_return_sequences"]):
             generated_tokens_logprobs = []
-            generated_tokens_top_logprobs_dicts = []
             for i in range(len(sequences[completion_id]) - len(encoded_input.input_ids[0])):
                 logprobs = torch.nn.functional.log_softmax(scores[i][completion_id], dim=0)
-                # Get top tokens in terms of log probability.
-                topk_logprobs = torch.topk(logprobs, k=top_k_per_token)
-                with self.wrapped_tokenizer as tokenizer:
-                    generated_tokens_top_logprobs_dicts.append(
-                        {
-                            tokenizer.convert_ids_to_tokens(k.item()): v.item()
-                            for (k, v) in zip(topk_logprobs.indices, topk_logprobs.values)
-                        }
-                    )
                 # Get log probability of chosen token.
                 j = i + len(encoded_input.input_ids[0])
                 generated_tokens_logprobs.append(logprobs[sequences[completion_id][j]].item())
             all_generated_tokens_logprobs.append(generated_tokens_logprobs)
-            all_generated_tokens_top_logprobs_dicts.append(generated_tokens_top_logprobs_dicts)
 
         # Remove prompt from the start of each sequence if echo_prompt is False.
         if not raw_request["echo_prompt"]:
@@ -162,17 +148,15 @@ class HuggingFaceServer:
             all_decoded_text = tokenizer.batch_decode(sequences)
 
         completions = []
-        for decoded_text, tokens, generated_tokens_logprobs, generated_tokens_top_logprobs_dicts in zip(
-            all_decoded_text, all_tokens, all_generated_tokens_logprobs, all_generated_tokens_top_logprobs_dicts
+        for decoded_text, tokens, generated_tokens_logprobs in zip(
+            all_decoded_text, all_tokens, all_generated_tokens_logprobs
         ):
             completions.append(
                 {
                     "text": decoded_text,
                     "tokens": tokens,
                     "logprobs": generated_tokens_logprobs,
-                    "top_logprobs_dicts": generated_tokens_top_logprobs_dicts,
                     "prompt_logprobs": prompt_tokens_logprobs,
-                    "prompt_top_logprobs_dicts": prompt_tokens_top_logprobs_dicts,
                 }
             )
 
@@ -240,7 +224,7 @@ class HuggingFaceClient(CachingClient):
         if request.embedding:
             return EMBEDDING_UNAVAILABLE_REQUEST_RESULT
 
-        raw_request = {
+        raw_request: HuggingFaceRequest = {
             "engine": request.model_engine,
             "prompt": request.prompt,
             "temperature": 1e-7 if request.temperature == 0 else request.temperature,
@@ -252,20 +236,18 @@ class HuggingFaceClient(CachingClient):
             "stop_sequences": request.stop_sequences,
         }
 
-        pretrained_model_name_or_path: str
-        if self._pretrained_model_name_or_path:
-            pretrained_model_name_or_path = self._pretrained_model_name_or_path
-        else:
-            pretrained_model_name_or_path = resolve_alias(request.model_deployment)
+        pretrained_model_name_or_path = (
+            self._pretrained_model_name_or_path if self._pretrained_model_name_or_path else request.model
+        )
         huggingface_model: HuggingFaceServer = HuggingFaceServerFactory.get_server(
-            helm_model_name=request.model_deployment,
+            helm_model_name=request.model,
             pretrained_model_name_or_path=pretrained_model_name_or_path,
             **self._kwargs,
         )
 
         try:
 
-            def do_it():
+            def do_it() -> Dict[str, Any]:
                 return huggingface_model.serve_request(raw_request)
 
             cache_key = CachingClient.make_cache_key(raw_request, request)
@@ -282,29 +264,26 @@ class HuggingFaceClient(CachingClient):
             if request.echo_prompt:
                 # Add prompt to list of generated tokens.
                 generated_tokens = raw_completion["tokens"][response["input_length"] :]
-                if raw_completion.get("prompt_logprobs") and raw_completion.get("prompt_top_logprobs_dicts"):
-                    for token_text, logprob, top_logprobs_dict in zip(
+                if raw_completion.get("prompt_logprobs"):
+                    for token_text, logprob in zip(
                         raw_completion["tokens"][: response["input_length"]],
                         raw_completion["prompt_logprobs"][: response["input_length"]],
-                        raw_completion["prompt_top_logprobs_dicts"][: response["input_length"]],
                     ):
-                        tokens.append(Token(text=token_text, logprob=logprob, top_logprobs=top_logprobs_dict))
+                        tokens.append(Token(text=token_text, logprob=logprob))
                         sequence_logprob += logprob
                 else:
                     for token_text in raw_completion["tokens"][: response["input_length"]]:
-                        tokens.append(Token(text=token_text, logprob=0.0, top_logprobs={}))
+                        tokens.append(Token(text=token_text, logprob=0.0))
 
             else:
                 generated_tokens = raw_completion["tokens"]
 
             # Compute logprob for the entire sequence.
-            for token_text, logprob, top_logprobs_dict in zip(
-                generated_tokens, raw_completion["logprobs"], raw_completion["top_logprobs_dicts"]
-            ):
-                tokens.append(Token(text=token_text, logprob=logprob, top_logprobs=top_logprobs_dict))
+            for token_text, logprob in zip(generated_tokens, raw_completion["logprobs"]):
+                tokens.append(Token(text=token_text, logprob=logprob))
                 sequence_logprob += logprob
 
-            completion = Sequence(text=raw_completion["text"], logprob=sequence_logprob, tokens=tokens)
+            completion = GeneratedOutput(text=raw_completion["text"], logprob=sequence_logprob, tokens=tokens)
             completion = truncate_sequence(completion, request)
             completions.append(completion)
 

helm/clients/image_generation/adobe_vision_client.py

@@ -0,0 +1,78 @@
+from typing import List, Dict
+
+from helm.common.cache import Cache, CacheConfig
+from helm.common.request import Request, RequestResult, GeneratedOutput
+from helm.common.tokenization_request import (
+    TokenizationRequest,
+    TokenizationRequestResult,
+    DecodeRequest,
+    DecodeRequestResult,
+)
+from helm.clients.client import Client, CachingClient
+from .image_generation_client_utils import get_single_image_multimedia_object
+
+
+class AdobeVisionClient(Client):
+    """
+    Client for Adobe vision models. Offline eval only.
+    """
+
+    SUPPORTED_MODELS: List[str] = ["giga-gan", "firefly"]
+
+    @staticmethod
+    def convert_to_raw_request(request: Request) -> Dict:
+        # Use default hyperparameters for everything else
+        raw_request: Dict = {
+            "request_type": "image-model-inference",
+            "model": request.model_engine,
+            "prompt": request.prompt,
+            "n": request.num_completions,
+        }
+        if request.random is not None:
+            raw_request["random"] = request.random
+        return raw_request
+
+    def __init__(self, cache_config: CacheConfig):
+        self._cache = Cache(cache_config)
+        self._promptist_model = None
+        self._promptist_tokenizer = None
+
+    def make_request(self, request: Request) -> RequestResult:
+        if request.model_engine not in self.SUPPORTED_MODELS:
+            raise ValueError(f"Unsupported model: {request.model_engine}")
+
+        raw_request = AdobeVisionClient.convert_to_raw_request(request)
+        raw_request.pop("random", None)
+        cache_key = CachingClient.make_cache_key(raw_request, request)
+
+        try:
+
+            def fail():
+                raise RuntimeError(
+                    f"The result has not been uploaded to the cache for the following request: {cache_key}"
+                )
+
+            response, cached = self._cache.get(cache_key, fail)
+        except RuntimeError as e:
+            error: str = f"Adobe Vision Client error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = [
+            GeneratedOutput(
+                text="", logprob=0, tokens=[], multimodal_content=get_single_image_multimedia_object(file_path)
+            )
+            for file_path in response["images"]
+        ]
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
+        raise NotImplementedError("This client does not support tokenizing.")
+
+    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
+        raise NotImplementedError("This client does not support decoding.")

helm/clients/image_generation/aleph_alpha_image_generation_client.py

@@ -0,0 +1,98 @@
+from typing import List, Dict
+
+from helm.common.cache import Cache, CacheConfig
+from helm.common.request import Request, RequestResult, GeneratedOutput
+from helm.common.tokenization_request import (
+    TokenizationRequest,
+    TokenizationRequestResult,
+    DecodeRequest,
+    DecodeRequestResult,
+)
+from helm.clients.client import Client, CachingClient
+from .image_generation_client_utils import get_single_image_multimedia_object
+
+
+class AlephAlphaImageGenerationClient(Client):
+    """
+    Client for Aleph Alpha vision models. Offline eval only.
+    """
+
+    DEFAULT_IMAGE_HEIGHT: int = 512
+    DEFAULT_IMAGE_WIDTH: int = 512
+
+    DEFAULT_GUIDANCE_SCALE: float = 7.5
+    DEFAULT_STEPS: int = 50
+
+    @staticmethod
+    def convert_to_raw_request(request: Request) -> Dict:
+        raw_request: Dict = {
+            "request_type": "image-model-inference",
+            "model": request.model_engine,
+            "prompt": request.prompt,
+            "n": request.num_completions,
+            "guidance_scale": AlephAlphaImageGenerationClient.DEFAULT_GUIDANCE_SCALE,
+            "steps": AlephAlphaImageGenerationClient.DEFAULT_STEPS,
+            "width": AlephAlphaImageGenerationClient.DEFAULT_IMAGE_WIDTH,
+            "height": AlephAlphaImageGenerationClient.DEFAULT_IMAGE_HEIGHT,
+        }
+        if request.random is not None:
+            raw_request["random"] = request.random
+
+        assert request.image_generation_parameters is not None
+        if request.image_generation_parameters.guidance_scale is not None:
+            raw_request["guidance_scale"] = request.image_generation_parameters.guidance_scale
+        if request.image_generation_parameters.diffusion_denoising_steps is not None:
+            raw_request["steps"] = request.image_generation_parameters.diffusion_denoising_steps
+        if (
+            request.image_generation_parameters.output_image_width is not None
+            and request.image_generation_parameters.output_image_height is not None
+        ):
+            raw_request["width"] = request.image_generation_parameters.output_image_width
+            raw_request["height"] = request.image_generation_parameters.output_image_height
+
+        return raw_request
+
+    def __init__(self, cache_config: CacheConfig):
+        self._cache = Cache(cache_config)
+        self._promptist_model = None
+        self._promptist_tokenizer = None
+
+    def make_request(self, request: Request) -> RequestResult:
+        if request.model_engine != "m-vader":
+            raise ValueError(f"Unsupported model: {request.model_engine}")
+
+        raw_request = AlephAlphaImageGenerationClient.convert_to_raw_request(request)
+        raw_request.pop("random", None)
+        cache_key = CachingClient.make_cache_key(raw_request, request)
+
+        try:
+
+            def fail():
+                raise RuntimeError(
+                    f"The result has not been uploaded to the cache for the following request: {cache_key}"
+                )
+
+            response, cached = self._cache.get(cache_key, fail)
+        except RuntimeError as e:
+            error: str = f"AlephAlphaVisionClient error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = [
+            GeneratedOutput(
+                text="", logprob=0, tokens=[], multimodal_content=get_single_image_multimedia_object(file_path)
+            )
+            for file_path in response["images"]
+        ]
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
+        raise NotImplementedError("This client does not support tokenizing.")
+
+    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
+        raise NotImplementedError("This client does not support decoding.")

helm/clients/image_generation/cogview2/coglm_strategy.py

@@ -0,0 +1,96 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   coglm_strategy.py
+@Time    :   2021/10/08 22:22:42
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+# here put the import lib
+import os
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+
+def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-65504):
+    # This function has been mostly taken from huggingface conversational ai code at
+    # https://medium.com/huggingface/how-to-build-a-state-of-the-art-conversational-ai-with-transfer-learning-2d818ac26313
+
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        # convert to 1D
+        logits = logits.view(logits.size()[1]).contiguous()
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+        # going back to 2D
+        logits = logits.view(1, -1).contiguous()
+
+    return logits
+
+
+class CoglmStrategy:
+    def __init__(
+        self, invalid_slices=[], temperature=1.0, top_k=200, eps=1e-4, top_p=0.0, end_tokens=None, top_k_cluster=1.0
+    ):
+        self.invalid_slices = invalid_slices
+        self.temperature = temperature
+        self.topk = top_k
+        self.top_p = top_p
+        self.eps = eps
+        if end_tokens is None:
+            end_tokens = []
+        self.end_tokens = end_tokens
+        self._is_done = False
+        self.outlier_count_down = 5
+        self.cluster_labels = torch.tensor(
+            np.load(f"{os.path.dirname(os.path.abspath(__file__))}/cluster_label.npy"),
+            device="cuda" if torch.cuda.is_available() else "cpu",
+            dtype=torch.long,
+        )
+        self.top_k_cluster = top_k_cluster
+
+    @property
+    def is_done(self) -> bool:
+        return self._is_done
+
+    def forward(self, logits, tokens, mems, temperature=None):
+        if temperature is None:
+            temperature = self.temperature
+        logits = logits / temperature
+        for invalid_slice in self.invalid_slices:
+            logits[..., invalid_slice] = -65504
+
+        rprobs = F.softmax(logits.float(), dim=-1)
+        c = self.cluster_labels.expand(*rprobs.shape)
+        cprobs = torch.zeros(logits.shape[0], 500, device=logits.device).scatter_add_(1, c, rprobs)
+        best_scores, best_clusters = cprobs.topk(self.topk)
+        bz = logits.shape[0]
+        for i in range(bz):
+            best_scores[i] = best_scores[i]  # ** 0.2
+            selected_cluster = best_clusters[i][torch.multinomial(best_scores[i] / best_scores[i].sum(), num_samples=1)]
+            logits[i, self.cluster_labels != selected_cluster] = -65504
+
+        probs = F.softmax(logits.float() / self.top_k_cluster, dim=-1)  # float is essential, due to a bug in Pytorch
+        pred = torch.multinomial(probs, num_samples=1)
+
+        if pred.numel() == 1 and pred.item() in self.end_tokens:
+            self._is_done = True
+        tokens = torch.cat((tokens, pred.view(tokens.shape[0], 1)), dim=1)
+        return tokens, mems
+
+    def finalize(self, tokens, mems):
+        self._is_done = False
+        return tokens, mems

helm/clients/image_generation/cogview2/coglm_utils.py

@@ -0,0 +1,82 @@
+import torch
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    from SwissArmyTransformer.model import CachedAutoregressiveModel
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+def get_masks_and_position_ids_coglm(seq, context_length):
+    tokens = seq.unsqueeze(0)
+    attention_mask = torch.ones((1, len(seq), len(seq)), device=tokens.device)
+    attention_mask.tril_()
+    attention_mask[..., :context_length] = 1
+    attention_mask.unsqueeze_(1)
+    position_ids = torch.zeros(len(seq), device=tokens.device, dtype=torch.long)
+    torch.arange(0, context_length, out=position_ids[:context_length])
+    torch.arange(512, 512 + len(seq) - context_length, out=position_ids[context_length:])
+    position_ids = position_ids.unsqueeze(0)
+    return tokens, attention_mask, position_ids
+
+
+def get_recipe(name):
+    r = {
+        "attn_plus": 1.4,
+        "temp_all_gen": 1.15,
+        "topk_gen": 16,
+        "temp_cluster_gen": 1.0,
+        "temp_all_dsr": 1.5,
+        "topk_dsr": 100,
+        "temp_cluster_dsr": 0.89,
+        "temp_all_itersr": 1.3,
+        "topk_itersr": 16,
+        "query_template": "{}<start_of_image>",
+    }
+    if name == "none":
+        pass
+    elif name == "mainbody":
+        r["query_template"] = "{} 高清摄影 隔绝<start_of_image>"
+    elif name == "photo":
+        r["query_template"] = "{} 高清摄影<start_of_image>"
+    elif name == "flat":
+        r["query_template"] = "{} 平面风格<start_of_image>"
+        # r['attn_plus'] = 1.8
+        # r['temp_cluster_gen'] = 0.75
+        r["temp_all_gen"] = 1.1
+        r["topk_dsr"] = 5
+        r["temp_cluster_dsr"] = 0.4
+        r["temp_all_itersr"] = 1
+        r["topk_itersr"] = 5
+    elif name == "comics":
+        r["query_template"] = "{} 漫画 隔绝<start_of_image>"
+        r["topk_dsr"] = 5
+        r["temp_cluster_dsr"] = 0.4
+        r["temp_all_gen"] = 1.1
+        r["temp_all_itersr"] = 1
+        r["topk_itersr"] = 5
+    elif name == "oil":
+        r["query_template"] = "{} 油画风格<start_of_image>"
+        pass
+    elif name == "sketch":
+        r["query_template"] = "{} 素描风格<start_of_image>"
+        r["temp_all_gen"] = 1.1
+    elif name == "isometric":
+        r["query_template"] = "{} 等距矢量图<start_of_image>"
+        r["temp_all_gen"] = 1.1
+    elif name == "chinese":
+        r["query_template"] = "{} 水墨国画<start_of_image>"
+        r["temp_all_gen"] = 1.12
+    elif name == "watercolor":
+        r["query_template"] = "{} 水彩画风格<start_of_image>"
+    return r
+
+
+class InferenceModel(CachedAutoregressiveModel):
+    def final_forward(self, logits, **kwargs):
+        logits_parallel = logits
+        logits_parallel = torch.nn.functional.linear(
+            logits_parallel.float(), self.transformer.word_embeddings.weight[:20000].float()
+        )
+        return logits_parallel

helm/clients/image_generation/cogview2/sr_pipeline/__init__.py

@@ -0,0 +1,15 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   __init__.py
+@Time    :   2022/03/02 13:57:09
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+from .direct_sr import DirectSuperResolution
+from .iterative_sr import IterativeSuperResolution
+from .sr_group import SRGroup
+
+DirectSuperResolution
+IterativeSuperResolution
+SRGroup