Update app.py

app.py

@@ -7,6 +7,58 @@ from gradio.themes.utils import colors
 
 from transformers import pipeline, TextIteratorStreamer, AutoModelForCausalLM, AutoTokenizer
 
+SYSTEM_PROMPT = "You are a compliance assistant. Use the provided risk data to answer user questions. If a single risk object is given, provide a direct answer. If a list of risks is provided, summarize, compare, or analyze the collection as needed. Always base your response on the data provided."
+
+class HfModelWrapper:
+    def __init__(
+        self,
+        model_path="casperhansen/llama-3.3-70b-instruct-awq",
+        sys_prompt=SYSTEM_PROMPT,
+        adapter_path="artemisiaai/fine-tuned-adapter",
+    ):
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path, device_map="auto", quantization_config=quantization, device_map="auto"
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.sys_prompt = sys_prompt
+        self.adapter_path = adapter_path
+        self.model.load_adapter(self.adapter_path)
+        self.model.enable_adapters()
+
+    def build_prompt(self, user_msg, history):
+
+        inppt = []
+        inppt.append({"role": "system", "content": self.sys_prompt})
+        inppt += history
+        inppt.append({"role": "user", "content": user_msg})
+
+        prompt = self.tokenizer.apply_chat_template(
+            inppt,
+            tokenize=False,
+        )
+        return prompt
+
+    def generate(self, user_input, history):
+        input_text = self.build_prompt(user_input, history)
+        input_ids = self.tokenizer.encode(input_text, return_tensors="pt").to("cuda")
+
+        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True)
+
+        gen_kwargs = {
+            "inputs": input_ids,
+            "streamer": streamer,
+            "pad_token_id": self.tokenizer.eos_token_id,
+            "max_length": 8192,
+            "temperature": 0.1,
+            "top_p": 0.8,
+            "repetition_penalty": 1.25,
+        }
+
+        thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
+        thread.start()
+
+        return streamer
 
 # Custom theme colors based on brand standards
 class ArtemisiaTheme(Base):
@@ -431,11 +483,8 @@ paper_plane_svg = """<svg xmlns="http://www.w3.org/2000/svg" width="20" height="
   <path d="M22 2L15 22L11 13L2 9L22 2Z"/>
 </svg>"""
 
+wrapper = HfModelWrapper()
 
-# Pipeline loading
-#generator = pipeline("text-generation", model="openai-community/gpt2")
-tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
-model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
 
 # Mock data function for chatbot
 def send_message(message, history):
@@ -444,23 +493,10 @@ def send_message(message, history):
     history.append({"role": "user", "content": message})
     #history.append({"role": "assistant", "content": f"This is a response about: {message}"})
     #return history
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-    input_ids = tokenizer.encode(message, return_tensors="pt")
-    gen_kwargs = {
-        "inputs": input_ids,
-        "streamer": streamer,
-        "pad_token_id": tokenizer.eos_token_id,
-        "max_length": 8192,
-        "temperature": 0.1,
-        "top_p": 0.8,
-        "repetition_penalty": 1.25,
-    }
+    response_generator = wrapper.generate(user_input, history)
     partial = ""
-    thread = Thread(target=model.generate, kwargs=gen_kwargs)
-    thread.start()    
-    #for token in generator(message, max_new_tokens=200):
-    for t in streamer:
-        partial += t#token["generated_text"][len(message):]
+    for t in response_generator:
+        partial += t
         yield history + [{"role": "assistant", "content": partial}]