Create app.py

app.py

@@ -0,0 +1,85 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModel, pipeline
+import faiss
+import numpy as np
+import torch
+import os
+from PyPDF2 import PdfReader
+
+# ===== 嵌入模型 =====
+embed_model = AutoModel.from_pretrained(
+    "BAAI/bge-small-zh", trust_remote_code=True
+)
+embed_tokenizer = AutoTokenizer.from_pretrained(
+    "BAAI/bge-small-zh", trust_remote_code=True
+)
+
+def embed_text(text):
+    inputs = embed_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        embeddings = embed_model(**inputs).last_hidden_state[:, 0, :]
+    return embeddings[0].numpy()
+
+# ===== 生成模型（轻量LLM） =====
+generator = pipeline(
+    "text-generation",
+    model="Qwen/Qwen1.5-1.8B-Chat-GGML",
+    device=-1
+)
+
+# ===== 全局变量存储索引和文档 =====
+index = None
+docs = []
+
+# ===== 文件解析函数 =====
+def load_file(file_obj):
+    global index, docs
+    docs = []
+    text_data = ""
+
+    if file_obj.name.endswith(".pdf"):
+        reader = PdfReader(file_obj.name)
+        for page in reader.pages:
+            text_data += page.extract_text() + "\n"
+    elif file_obj.name.endswith(".txt"):
+        text_data = file_obj.read().decode("utf-8")
+    else:
+        return "仅支持 PDF 或 TXT 文件", None
+
+    # 切块
+    chunks = [text_data[i:i+500] for i in range(0, len(text_data), 500)]
+    docs = [{"text": chunk, "source": f"chunk_{i}"} for i, chunk in enumerate(chunks)]
+
+    # 向量化并建索引
+    doc_embeddings = np.array([embed_text(d["text"]) for d in docs])
+    index = faiss.IndexFlatL2(doc_embeddings.shape[1])
+    index.add(doc_embeddings)
+
+    return f"已加载 {len(docs)} 个文本块", None
+
+# ===== RAG 查询函数 =====
+def rag_query(query):
+    if index is None:
+        return "请先上传文件构建知识库"
+    q_emb = embed_text(query).reshape(1, -1)
+    D, I = index.search(q_emb, k=3)
+    retrieved = [docs[i]["text"] for i in I[0]]
+    context = "\n".join(retrieved)
+    prompt = f"已知信息：\n{context}\n\n问题：{query}\n请基于已知信息回答，并引用来源。"
+    result = generator(prompt, max_length=200, do_sample=False)
+    return result[0]["generated_text"]
+
+# ===== Gradio 界面 =====
+with gr.Blocks() as demo:
+    gr.Markdown("## 📚 轻量 RAG 原型（上传 PDF/TXT）")
+    with gr.Row():
+        file_input = gr.File(label="上传 PDF 或 TXT 文件")
+        load_btn = gr.Button("构建知识库")
+    status = gr.Textbox(label="状态")
+    query_input = gr.Textbox(label="输入你的问题")
+    answer_output = gr.Textbox(label="回答")
+    load_btn.click(load_file, inputs=file_input, outputs=status)
+    query_input.submit(rag_query, inputs=query_input, outputs=answer_output)
+
+if __name__ == "__main__":
+    demo.launch()