app.py

import gradio as gr
import fitz  # PyMuPDF
from langchain.llms import Ollama
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate
from langchain.docstore.document import Document

def extract_text_from_pdf(file):
    doc = fitz.open(stream=file.read(), filetype="pdf")
    text = ""
    for page in doc:
        text += page.get_text()
    return text

def answer_question(pdf_file, question):
    if not pdf_file or not question:
        return "Vui lòng tải file PDF và nhập câu hỏi."

    content = extract_text_from_pdf(pdf_file)

    # Khởi tạo model Ollama
    llm = Ollama(model="llama3")  # Bạn có thể dùng 'mistral', 'phi', v.v.

    # Prompt đơn giản
    prompt_template = PromptTemplate.from_template(
        "Trả lời câu hỏi sau dựa trên tài liệu:\n\n{context}\n\nCâu hỏi: {question}"
    )

    chain = load_qa_chain(llm, chain_type="stuff", prompt=prompt_template)

    # LangChain yêu cầu list Document
    docs = [Document(page_content=content)]

    result = chain.run(input_documents=docs, question=question)
    return result

gr.Interface(
    fn=answer_question,
    inputs=[
        gr.File(label="Tải lên PDF"),
        gr.Textbox(label="Câu hỏi về nội dung PDF")
    ],
    outputs="text",
    title="PDF Q&A với Ollama",
    description="Sử dụng mô hình Ollama (như llama3) để trả lời câu hỏi từ file PDF.",
).launch()