Update app.py

app.py

@@ -29,38 +29,49 @@ function refresh() {
 
 previous_sessions = []
 
-class ChatState():
-    __START_TURN_USER__ = "<start_of_turn>user\n"
-    __START_TURN_MODEL__ = "<start_of_turn>model\n"
-    __END_TURN__ = "<end_of_turn>\n"
-
-    def __init__(self, model, system=""):
+class ChatState:
+    def __init__(self, model, tokenizer, system=""):
         self.model = model
+        self.tokenizer = tokenizer
         self.system = system
         self.history = []
 
-    def add_to_history_as_user(self, message):
-        self.history.append(self.__START_TURN_USER__ + message + self.__END_TURN__)
-
-    def add_to_history_as_model(self, message):
-        self.history.append(self.__START_TURN_MODEL__ + message)
-
-    def get_history(self):
-        return "".join([*self.history])
+    def add_to_history(self, role, message):
+        self.history.append({"role": role, "content": message})
 
     def get_full_prompt(self):
-        prompt = self.get_history() + self.__START_TURN_MODEL__
-        if len(self.system) > 0:
-            prompt = self.system + "\n" + prompt
+        prompt = ""
+        if self.system:
+            prompt += f"System: {self.system}\n\n"
+        for message in self.history:
+            prompt += f"{message['role'].capitalize()}: {message['content']}\n"
+        prompt += "Model: "
         return prompt
 
     def send_message(self, message):
-        self.add_to_history_as_user(message)
+        self.add_to_history("user", message)
         prompt = self.get_full_prompt()
-        response = self.model.generate(prompt, max_length=2048)
-        result = response.replace(prompt, "")  # Extract only the new response
-        self.add_to_history_as_model(result)
-        return result
+        
+        inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=512,
+                num_return_sequences=1,
+                do_sample=True,
+                temperature=0.7
+            )
+        
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response.replace(prompt, "").strip()
+        
+        self.add_to_history("model", response)
+        return response
+
+# Initialize the ChatState
+chat_state = ChatState(model, tokenizer, system="You are a helpful AI assistant.")
+
 
 def post_process_output(prompt, result):
     answer = result.strip()
@@ -102,18 +113,15 @@ def inference(prompt):
     if reset_triggered:
         return "", ""
     
-    chat_state.send_message(prompt)  # Process the user's message
-
-    # Post-process the output from the model
-    formatted_output = post_process_output(chat_state.get_full_prompt(), chat_state.get_history())
+    result = chat_state.send_message(prompt)
     
     # Apply a bit of delay for a realistic response time
     time.sleep(1)
     
-    result = formatted_output
-    sessions = add_session(chat_state.get_history())
+    sessions = add_session(prompt)
     return result, sessions
 
+
 with gr.Blocks(theme='gradio/soft', css=custom_css) as demo:
     
     gr.Markdown("<center><h1>HydroSense LLM Demo</h1></center>")