Save ram

app.py

@@ -35,104 +35,15 @@ import tensorflow_text
 from dataclasses import dataclass
 
 import numpy as np
-import tensorflow as tf
 
 
-class Encoder(ABC):
-    @abstractmethod
-    def encode(self, texts: List[str]) -> np.array:
-      """
-        output dimension expected to be one dimension and normalized (unit vector)
-      """
-      ...
-
-
-class MUSEEncoder(Encoder):
-    def __init__(self, model_url: str = "https://tfhub.dev/google/universal-sentence-encoder-multilingual/3"):
-        self.embed = hub.load(model_url)
-
-    def encode(self, texts: List[str]) -> np.array:
-        embeds = self.embed(texts).numpy()
-        embeds = embeds / np.linalg.norm(embeds, axis=1).reshape(embeds.shape[0], -1)
-        return embeds
-
-
-
-
-
-@dataclass
-class SensitiveTopic:
-    name: str
-    respond_message: str
-    sensitivity: float = None # range from 0 to 1
-    demonstrations: List[str] = None
-    adhoc_embeded_demonstrations: np.array = None # dimension = [N_ADHOC, DIM]. Please kindly note that this suppose to 
-
-
-DEFAULT_SENSITIVITY = 0.4
-
-
-class SensitiveTopicProtector:
-    def __init__(
-        self,
-        sensitive_topics: List[SensitiveTopic],
-        encoder: Encoder = MUSEEncoder(),
-        default_sensitivity: float = DEFAULT_SENSITIVITY
-    ):
-        self.sensitive_topics = sensitive_topics
-        self.default_sensitivity = default_sensitivity
-        self.encoder = encoder
-        self.topic_embeddings = self._get_topic_embeddings()
-
-    def _get_topic_embeddings(self) -> Dict[str, List[np.array]]:
-        topic_embeddings = {}
-        for topic in self.sensitive_topics:
-            current_topic_embeddings = None
-            if topic.demonstrations is not None:
-                current_topic_embeddings = self.encoder.encode(texts=topic.demonstrations) if current_topic_embeddings is None \
-                    else np.concatenate((current_topic_embeddings, self.encoder.encode(texts=topic.demonstrations)), axis=0)
-            if topic.adhoc_embeded_demonstrations is not None:
-                current_topic_embeddings = topic.adhoc_embeded_demonstrations if current_topic_embeddings is None \
-                    else np.concatenate((current_topic_embeddings, topic.adhoc_embeded_demonstrations), axis=0)
-            topic_embeddings[topic.name] = current_topic_embeddings
-        return topic_embeddings
-
-    def filter(self, text: str) -> Tuple[bool, str]:
-        is_sensitive, respond_message = False, None
-        text_embedding = self.encoder.encode([text,])
-        for topic in self.sensitive_topics:
-            risk_scores = np.einsum('ik,jk->j', text_embedding, self.topic_embeddings[topic.name])
-            max_risk_score = np.max(risk_scores)
-            if topic.sensitivity:
-                if max_risk_score > (1.0 - topic.sensitivity):
-                    return True, topic.respond_message
-                continue
-            if max_risk_score > (1.0 - self.default_sensitivity):
-                return True, topic.respond_message
-        return is_sensitive, respond_message
-    
-    @classmethod
-    def fromRaw(cls, raw_sensitive_topics: List[Dict], encoder: Encoder = MUSEEncoder(), default_sensitivity: float = DEFAULT_SENSITIVITY):
-        sensitive_topics = [SensitiveTopic(**topic) for topic in raw_sensitive_topics]
-        return cls(sensitive_topics=sensitive_topics, encoder=encoder, default_sensitivity=default_sensitivity)
-
-
-
-f = open("sensitive_topics.pkl", "rb")
-sensitive_topics = pickle.load(f)
-f.close()
-
-guardian = SensitiveTopicProtector.fromRaw(sensitive_topics)
-
-
-name_model = "pythainlp/wangchanglm-7.5B-sft-en-8bit-sharded"
+name_model = "pythainlp/wangchanglm-7.5B-sft-en-sharded"
 model = AutoModelForCausalLM.from_pretrained(
     name_model, 
     device_map="auto", 
     torch_dtype=torch.bfloat16, 
     offload_folder="./", 
     low_cpu_mem_usage=True,
-    load_in_8bit=False
 )
 tokenizer = AutoTokenizer.from_pretrained("facebook/xglm-7.5B")
 
@@ -238,9 +149,6 @@ def gen_instruct(text,max_new_tokens=512,top_p=0.95,temperature=0.9,top_k=50):
     return tokenizer.decode(output_tokens[0][len(batch["input_ids"][0]):], skip_special_tokens=True)
 
 def gen_chatbot_old(text):
-    is_sensitive, respond_message = guardian.filter(text)
-    if is_sensitive:
-        return respond_message
 
     batch = tokenizer(text, return_tensors="pt")
     #context_tokens = tokenizer(text, add_special_tokens=False)['input_ids']
@@ -278,9 +186,6 @@ def instruct_generate(
     temperature: float = 0.1,
     top_p: float = 0.75,
 ):
-    is_sensitive, respond_message = guardian.filter(instruct)
-    if is_sensitive:
-        return respond_message
 
     if input == 'none' or len(input)<2:
         prompt = PROMPT_DICT['prompt_no_input'].format_map(
@@ -391,11 +296,7 @@ with gr.Blocks(height=900) as demo:
               x=sumbit_data(save="chat",prompt=_bot,vote=3,feedback=feedback)
             return {feedback_chatbot_ok: gr.update(visible=True),feedback_chatbot_box: gr.update(visible=False)}
         def user(user_message, history):
-            is_sensitive, respond_message = guardian.filter(user_message)
-            if is_sensitive:
-                bot_message = respond_message
-            else:
-                bot_message = chatgpt_chain.predict(human_input=user_message)
+            bot_message = chatgpt_chain.predict(human_input=user_message)
             history.append((user_message, bot_message))
             return "", history,gr.update(visible=True)
         def reset():