Create app.py

app.py

@@ -0,0 +1,61 @@
+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"
+    return pdf_Text
+
+def generate_response(pdf, query):
+    pdf_Text = extract_text_from_pdf(pdf)
+    input_text = prompt.format(pdf_Text, query, " ")
+    input_ids = tokenizer(input_text, return_tensors="pt")
+
+    if torch.cuda.is_available():
+        input_ids = input_ids.to("cuda")
+
+    outputs = model.generate(
+        **input_ids,
+        max_new_tokens=500,  # Limit the number of tokens generated
+        no_repeat_ngram_size=5,  # Prevent repetition of 5-gram phrases
+    )
+    return tokenizer.decode(outputs[0])
+
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=[gr.inputs.File(label="Upload PDF"), gr.inputs.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()