Create app.py

app.py

@@ -0,0 +1,56 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from PyPDF2 import PdfReader
+import torch
+
+# Load IBM Granite model (use a smaller one if needed, e.g., granite-3.0-3b-instruct)
+model_name = "ibm-granite/granite-3.0-8b-instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+
+# Global variable for PDF context
+pdf_context = ""
+
+def upload_pdf(file):
+    global pdf_context
+    if file is None:
+        return "No file uploaded."
+    reader = PdfReader(file)
+    pdf_context = ""
+    for page in reader.pages:
+        pdf_context += page.extract_text() + "\n"
+    return "PDF uploaded and text extracted successfully!"
+
+def chat(message, history):
+    # Build messages with history and PDF context
+    messages = [{"role": "system", "content": f"You are a helpful assistant. Answer based on this context: {pdf_context}"}]
+    for user_msg, assistant_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": assistant_msg})
+    messages.append({"role": "user", "content": message})
+    
+    # Apply chat template and generate
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=300, temperature=0.7)
+    response = tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+    return response
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Basic PDF Q&A Chat with IBM Granite")
+    
+    with gr.Row():
+        pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
+        upload_btn = gr.Button("Upload PDF")
+    
+    status = gr.Textbox(label="Status")
+    
+    chat_interface = gr.ChatInterface(
+        fn=chat,
+        title="Ask questions about the PDF"
+    )
+    
+    upload_btn.click(upload_pdf, inputs=pdf_input, outputs=status)
+
+demo.launch()