Update app.py

app.py

@@ -1,17 +1,11 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
 from pathlib import Path
-from typing import List
+from transformers import RAGTokenForingModel, AutoTokenizer, AutoModelForCausalLM
 from pdfplumber import open as open_pdf
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+from typing import List
 
 # Load the PDF file
-pdf_path = Path("apexcustoms.pdf")
+pdf_path = Path("path/to/your/pdf/file.pdf")
 with open_pdf(pdf_path) as pdf:
     text = "\n".join(page.extract_text() for page in pdf.pages)
 
@@ -19,9 +13,13 @@ with open_pdf(pdf_path) as pdf:
 chunk_size = 1000  # Adjust this value based on your needs
 text_chunks: List[str] = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
 
-# Load the AutoModelForCausalLM and tokenizer
-model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-tpu-automotive")
-tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-tpu-automotive")
+# Load the RAG model and tokenizer for retrieval
+rag_tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
+rag_model = RAGTokenForingModel.from_pretrained("facebook/rag-token-nq")
+
+# Load the DialoGPT model and tokenizer for generation
+dialogpt_tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+dialogpt_model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 
 def respond(
     message,
@@ -43,15 +41,16 @@ def respond(
 
     response = ""
 
-    # Pass relevant chunks as context
-    relevant_chunks = [chunk for chunk in text_chunks if message.lower() in chunk.lower()]
-    context = "\n".join(relevant_chunks)
+    # Retrieve relevant chunks using the RAG model
+    rag_input_ids = rag_tokenizer(message, return_tensors="pt").input_ids
+    rag_output = rag_model(rag_input_ids, text_chunks, return_retrieved_inputs=True)
+    retrieved_text = rag_output.retrieved_inputs
 
-    # Encode the context and user's message
-    input_ids = tokenizer.encode(context + "\n\n" + message, return_tensors="pt")
+    # Encode the context and user's message for DialoGPT
+    input_ids = dialogpt_tokenizer.encode(retrieved_text + "\n\n" + message, return_tensors="pt")
 
-    # Generate the response using the AutoModelForCausalLM
-    output = model.generate(
+    # Generate the response using the DialoGPT model
+    output = dialogpt_model.generate(
         input_ids,
         max_length=max_tokens,
         num_beams=num_beams,
@@ -59,7 +58,7 @@ def respond(
         early_stopping=True
     )
 
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    response = dialogpt_tokenizer.decode(output[0], skip_special_tokens=True)
 
     yield response