Update app.py

app.py

@@ -1,9 +1,10 @@
 import os
 import torch
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from PyPDF2 import PdfReader
 import gradio as gr
-from datasets import Dataset
+from datasets import Dataset, load_from_disk, save_to_disk
+import faiss
 
 # Extract text from PDF
 def extract_text_from_pdf(pdf_path):
@@ -16,7 +17,8 @@ def extract_text_from_pdf(pdf_path):
 
 # Load model and tokenizer
 model_name = "scb10x/llama-3-typhoon-v1.5x-8b-instruct"
-tokenizer = RagTokenizer.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Extract text from the provided PDF
 pdf_text = extract_text_from_pdf("TOPF 2564.pdf")  # Updated path
@@ -29,23 +31,27 @@ dataset = Dataset.from_list(passages)
 dataset_path = "./rag_document_dataset"
 index_path = "./rag_document_index"
 
+# Ensure the directory exists
+os.makedirs(dataset_path, exist_ok=True)
+os.makedirs(index_path, exist_ok=True)
+
 # Save the dataset to disk and create an index
 dataset.save_to_disk(dataset_path)
 dataset.load_from_disk(dataset_path).add_faiss_index(column="text").save(index_path)
 
-# Load the retriever with the custom dataset and index
-retriever = RagRetriever.from_pretrained(
-    model_name,
-    index_name="custom",
-    passages_path=dataset_path,
-    index_path=index_path
-)
-
-model = RagSequenceForGeneration.from_pretrained(model_name, retriever=retriever)
+# Custom retriever
+def retrieve(query):
+    # Use FAISS index to retrieve relevant passages
+    query_embedding = tokenizer(query, return_tensors="pt")["input_ids"]
+    # Perform retrieval (this is a placeholder, actual retrieval code will be more complex)
+    # retrieved_passages = faiss_search(query_embedding)
+    retrieved_passages = " ".join([passage['text'] for passage in passages])  # Simplified for demo
+    return retrieved_passages
 
 # Define the chat function
 def answer_question(question, context):
-    inputs = tokenizer(question, return_tensors="pt")
+    retrieved_context = retrieve(question)
+    inputs = tokenizer(question + " " + retrieved_context, return_tensors="pt")
     input_ids = inputs["input_ids"]
     attention_mask = inputs["attention_mask"]