Update app.py

app.py

@@ -1,9 +1,77 @@
+import gradio as gr
+from PyPDF2 import PdfReader
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("himmeow/vi-gemma-2b-RAG")
+model = AutoModelForCausalLM.from_pretrained(
+    "himmeow/vi-gemma-2b-RAG",
+    device_map="auto",
+    torch_dtype=torch.bfloat16
+)
+
+if torch.cuda.is_available():
+    model.to("cuda")
+
+# Define the prompt format for the model
+prompt = """
+### Instruction and Input:
+Based on the following context/document:
+{}
+Please answer the question: {}
+
+### Response:
+{}
+"""
+
 def extract_text_from_pdf(pdf):
     pdf_Text = ""
     reader = PdfReader(pdf)
     for page_num in range(len(reader.pages)):
         page = reader.pages[page_num]
         text = page.extract_text()
-        pdf_Text += text + "\n"
-    print("Extracted Text:\n", pdf_Text)  # Add this line to debug
+        if text:
+            pdf_Text += text + "\n"
+    if not pdf_Text.strip():
+        pdf_Text = "The PDF contains no extractable text."
+    print("Extracted Text:\n", pdf_Text)  # Debugging statement
     return pdf_Text
+
+def generate_response(pdf, query):
+    pdf_Text = extract_text_from_pdf(pdf)
+    if not pdf_Text.strip():
+        return "The PDF appears to be empty or unreadable."
+    
+    input_text = prompt.format(pdf_Text, query, " ")
+    print("Input Text for Model:\n", input_text)  # Debugging statement
+
+    input_ids = tokenizer(input_text, return_tensors="pt")
+
+    if torch.cuda.is_available():
+        input_ids = input_ids.to("cuda")
+
+    try:
+        outputs = model.generate(
+            **input_ids,
+            max_new_tokens=500,
+            no_repeat_ngram_size=5,
+        )
+        response = tokenizer.decode(outputs[0])
+    except Exception as e:
+        response = "An error occurred while generating the response."
+        print("Error:", e)
+    
+    print("Generated Response:\n", response)  # Debugging statement
+    return response
+
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=[gr.File(label="Upload PDF"), gr.Textbox(label="Ask a question")],
+    outputs="text",
+    title="PDF Question Answering with vi-gemma-2b-RAG",
+    description="Upload a PDF and ask a question based on its content. The model will generate a response."
+)
+
+iface.launch()