Update app.py

app.py

@@ -19,8 +19,8 @@ with open(pdf_path, 'rb') as f:
 split_pattern = r'\n\n'
 doc_chunks = re.split(split_pattern, pdf_text)
 
-# Create the retriever input
-corpus = [{"text": chunk} for chunk in doc_chunks]
+# Preprocess the corpus
+corpus = rag_tokenizer(doc_chunks, return_tensors="pt", padding=True, truncation=True).input_ids
 
 """
 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
@@ -46,11 +46,20 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    # Tokenize the input and retrieve relevant context from the PDF
+    # Tokenize the input
     inputs = rag_tokenizer(message, return_tensors="pt")
-    inputs.update({"corpus": corpus})
     input_ids = inputs.pop("input_ids")
-    output_ids = rag_model.generate(**inputs, max_length=max_tokens, temperature=temperature, top_p=top_p, num_beams=2)
+
+    # Generate with the RAG model
+    output_ids = rag_model.generate(
+        input_ids=input_ids,
+        context_input_ids=corpus,
+        max_length=max_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=0,
+        num_beams=2,
+    )
     retrieved_context = rag_tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
 
     response = ""