Update client.py

client.py

@@ -1,15 +1,18 @@
-from huggingface_hub import InferenceClient
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 from init import ACCESS_TOKEN, SYSTEM_PROMPT
 from utils import extract_sql, is_sql
 from database import execute
-import os
-
-client = InferenceClient(api_key=os.environ.get('HF_TOKEN'))
 
+# Load the model and tokenizer
+model_name = "Qwen/Qwen2.5-3B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
 def respond(message, history, system_message, max_tokens, temperature, top_p):
     messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    # Xử lý lịch sử chat
+    
+    # Process chat history
     for val in history:
         if val[0]:
             messages.append({"role": "user", "content": val[0]})
@@ -18,21 +21,23 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
 
     messages.append({"role": "user", "content": message})
 
-    # Tạo response đầu tiên
-    response = ""
-    for message in client.chat.completions.create(
-        model="Qwen/Qwen2.5-3B-Instruct",
-        max_tokens=max_tokens,
-        stream=True,
+    # Tokenize input
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
+
+    # Generate response
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-        messages=messages,
-    ):
-        token = message.choices[0].delta.content
-        response += token
-        yield response
+        do_sample=True
+    )
 
-    # Xử lý logic SQL và retry
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    yield response
+
+    # SQL Processing and Retry Logic
     if is_sql(response):
         sql_query = extract_sql(response)
         max_attempts = 3
@@ -48,59 +53,52 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
                 last_error = str(e)
                 attempts += 1
                 if attempts < max_attempts:
-                    # Thêm thông tin lỗi vào context và yêu cầu mô hình hỏi lại người dùng
-                    clarification_prompt = f"""Tôi gặp lỗi khi thực hiện truy vấn SQL: {last_error}
-                    Bạn có thể cung cấp thêm thông tin hoặc chỉnh sửa câu hỏi để tôi có thể sửa truy vấn không?"""
+                    clarification_prompt = f"Tôi gặp lỗi khi thực hiện truy vấn SQL: {last_error}\nBạn có thể chỉnh sửa câu hỏi hoặc cung cấp thêm thông tin không?"
                     messages += [
                         {"role": "assistant", "content": response},
                         {"role": "user", "content": clarification_prompt},
                     ]
 
-                    # Tạo response yêu cầu thông tin thêm
-                    response = ""
-                    for message in client.chat.completions.create(
-                        model="Qwen/Qwen2.5-3B-Instruct",
-                        max_tokens=max_tokens,
-                        stream=True,
+                    # Tokenize clarification prompt
+                    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+                    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
+
+                    # Generate new response
+                    output_ids = model.generate(
+                        input_ids,
+                        max_new_tokens=max_tokens,
                         temperature=temperature,
                         top_p=top_p,
-                        messages=messages,
-                    ):
-                        token = message.choices[0].delta.content
-                        response += token
-                        yield response
+                        do_sample=True
+                    )
+
+                    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+                    yield response
 
-                    # Nếu mô hình cung cấp SQL mới, tiếp tục thử
                     if is_sql(response):
                         sql_query = extract_sql(response)
                 else:
-                    # Nếu sau 3 lần vẫn lỗi, tiếp tục hỏi lại người dùng thay vì in lỗi
-                    retry_prompt = f"""Tôi đã thử {max_attempts} lần nhưng vẫn gặp lỗi: {last_error}
-                    Bạn có thể cung cấp thêm chi tiết về dữ liệu cần truy vấn không?"""
-                    messages.append({"role": "assistant", "content": retry_prompt})
+                    retry_prompt = f"Tôi đã thử {max_attempts} lần nhưng vẫn gặp lỗi: {last_error}\nBạn có thể cung cấp thêm chi ti���t về dữ liệu cần truy vấn không?"
                     yield retry_prompt
                     return
 
-        # Nếu thực hiện truy vấn thành công
         if sql_result is not None:
-            reformulation_prompt = f"""Kết quả truy vấn SQL:
-            {sql_result}
-            Hãy tóm tắt kết quả thành phản hồi tự nhiên cho người dùng."""
+            reformulation_prompt = f"Kết quả truy vấn SQL:\n{sql_result}\nHãy tóm tắt kết quả thành phản hồi tự nhiên."
             messages += [
                 {"role": "assistant", "content": response},
                 {"role": "user", "content": reformulation_prompt},
             ]
 
-            # Tạo response tóm tắt
-            reformulated_response = ""
-            for message in client.chat.completions.create(
-                model="Qwen/Qwen2.5-3B-Instruct",
-                max_tokens=512,
-                stream=True,
+            input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)
+
+            output_ids = model.generate(
+                input_ids,
+                max_new_tokens=512,
                 temperature=temperature,
                 top_p=top_p,
-                messages=messages,
-            ):
-                token = message.choices[0].delta.content
-                reformulated_response += token
-                yield reformulated_response
+                do_sample=True
+            )
+
+            reformulated_response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+            yield reformulated_response