Update app.py

app.py

@@ -1,21 +1,70 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
+import torch
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from PyPDF2 import PdfReader
 import gradio as gr
+from datasets import Dataset
+
+# Extract text from PDF
+def extract_text_from_pdf(pdf_path):
+    text = ""
+    with open(pdf_path, "rb") as f:
+        reader = PdfReader(f)
+        for page in reader.pages:
+            text += page.extract_text()
+    return text
+
+# Load model and tokenizer
+model_name = "scb10x/llama-3-typhoon-v1.5x-8b-instruct"
+tokenizer = RagTokenizer.from_pretrained(model_name)
+
+# Extract text from the provided PDF
+pdf_text = extract_text_from_pdf("TOPF 2564.pdf")  # Updated path
+passages = [{"title": "", "text": line} for line in pdf_text.split('\n') if line.strip()]
+
+# Create a Dataset
+dataset = Dataset.from_list(passages)
+
+# Save the dataset and create an index in the current working directory
+dataset_path = "./rag_document_dataset"
+index_path = "./rag_document_index"
+
+# 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)
 
-model_name = "scb10x/llama-3-typhoon-v1.5x-70b-instruct-awq"
-token = os.getenv("HF_TOKEN")
+# Define the chat function
+def answer_question(question, context):
+    inputs = tokenizer(question, return_tensors="pt")
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
 
-# Remove these lines
-# device = torch.device("cuda")
-# torch.cuda.set_device(0)
+    # Generate the answer
+    outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return answer
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
-model = AutoModelForCausalLM.from_pretrained(model_name, token=token)
+# Gradio interface setup
+def ask(question):
+    return answer_question(question, pdf_text)
 
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs.input_ids, max_length=50)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+demo = gr.Interface(
+    fn=ask,
+    inputs=gr.inputs.Textbox(lines=2, placeholder="Ask something..."),
+    outputs="text",
+    title="Document QA with RAG",
+    description="Ask questions based on the provided document."
+)
 
-gr.Interface(fn=generate_text, inputs="text", outputs="text").launch()
+if __name__ == "__main__":
+    demo.launch()