Update app.py

app.py

@@ -3,6 +3,7 @@ from huggingface_hub import InferenceClient
 from pathlib import Path
 from typing import List
 from pdfplumber import open as open_pdf
+from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, 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
@@ -10,7 +11,7 @@ For more information on `huggingface_hub` Inference API support, please check th
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 # 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)
 
@@ -18,6 +19,11 @@ 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 VisionEncoderDecoderModel and tokenizer
+model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -40,19 +46,26 @@ def respond(
 
     # Pass relevant chunks as context
     relevant_chunks = [chunk for chunk in text_chunks if message.lower() in chunk.lower()]
+    context = "\n".join(relevant_chunks)
+
+    # Encode the context and user's message
+    inputs = feature_extractor(context, return_tensors="pt")
+    query = tokenizer(message, return_tensors="pt")
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
+    # Generate the response using the VisionEncoderDecoderModel
+    output_ids = model.generate(
+        **inputs,
+        **query,
+        max_length=max_tokens,
         temperature=temperature,
         top_p=top_p,
-        files={"context": "\n".join(relevant_chunks)},  # Pass relevant chunks as context
-    ):
-        token = message.choices[0].delta.content
+        num_beams=4,
+        early_stopping=True
+    )
+
+    response = tokenizer.decode(output_ids.squeeze(), skip_special_tokens=True)
 
-        response += token
-        yield response
+    yield response
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface