Update app.py

app.py

@@ -1,30 +1,13 @@
 import gradio as gr
 import pdfplumber
 import os
-from huggingface_hub import HfApi, InferenceClient
-import time
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Thông tin Dataset và Hugging Face Token của bạn
-DATASET_REPO = "blackanAK06/chatbox-dataset"  # Tên dataset của bạn trên Hugging Face
-API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")  # Đặt API Token của bạn từ biến môi trường
-
-# Khởi tạo Hugging Face API và InferenceClient với mô hình GPT-Neo
-api = HfApi()
-client = InferenceClient(model="EleutherAI/gpt-neo-2.7B", token=API_TOKEN)  # Khởi tạo InferenceClient với mô hình cụ thể
-
-# Hàm để tải lên dataset
-def upload_to_dataset(file_path, filename):
-    try:
-        api.upload_file(
-            path_or_fileobj=file_path,
-            path_in_repo=f"files/{filename}",
-            repo_id=DATASET_REPO,
-            repo_type="dataset",
-            token=API_TOKEN,
-        )
-        return f"Successfully uploaded {filename} to the dataset."
-    except Exception as e:
-        return f"Error uploading file: {e}"
+# Tải mô hình từ Hugging Face Model Hub
+model_name = "antphb/DS-Chatbox-facebook-xglm-564M-V4-FT"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Hàm để trích xuất văn bản từ tệp PDF
 def extract_text_from_pdf(file_path):
@@ -40,17 +23,12 @@ def extract_text_from_pdf(file_path):
     return text
 
 # Hàm để xử lý câu hỏi và nội dung tệp sau khi người dùng yêu cầu
-def respond_with_file(filepath, message, max_retries=3):
+def respond_with_file(filepath, message):
     if not message:
         return "Please enter a question."
 
     filename = os.path.basename(filepath)
 
-    # Tải tệp lên dataset
-    upload_message = upload_to_dataset(filepath, filename)
-    if "Error" in upload_message:
-        return upload_message
-
     # Trích xuất nội dung từ tệp đã tải lên
     if filename.endswith(".pdf"):
         file_content = extract_text_from_pdf(filepath)
@@ -68,21 +46,15 @@ def respond_with_file(filepath, message, max_retries=3):
     # Kết hợp nội dung tệp với câu hỏi để mô hình trả lời
     full_message = f"Here is the content of the file:\n{file_content}\n\nUser question: {message}"
 
-    # Cơ chế thử lại nếu gặp lỗi
-    for attempt in range(max_retries):
-        try:
-            # Gửi yêu cầu tới mô hình GPT-Neo để tạo ra câu trả lời
-            response = client.text_generation(full_message, max_new_tokens=150, temperature=0.7)
-            return f"{upload_message}\n\nResponse:\n{response}"
-        except Exception as e:
-            if "CANCELLED" in str(e):
-                time.sleep(2)  # Chờ 2 giây trước khi thử lại
-                continue
-            else:
-                return f"Error communicating with the model: {e}"
-
-    # Nếu hết số lần thử mà vẫn lỗi
-    return "Request failed after several attempts. Please try again later."
+    # Sử dụng mô hình để tạo phản hồi
+    try:
+        inputs = tokenizer(full_message, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=150, temperature=0.7)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return f"Response:\n{generated_text}"
+    except Exception as e:
+        return f"Error generating response from the model: {e}"
 
 # Giao diện Gradio
 def interface():
@@ -99,8 +71,8 @@ def interface():
         outputs="text",
         live=False,
         flagging_mode="never",  # Thay thế allow_flagging bằng flagging_mode
-        title="File-based Chatbot with Dataset Storage",
-        description="Upload a text file or PDF, save it to a dataset, and ask questions based on its content."
+        title="File-based Chatbot with Direct Model Access",
+        description="Upload a text file or PDF, and ask questions based on its content using a transformer model."
     )
     return demo