Spaces:
Running
Running
Update flask_app.py
Browse files- flask_app.py +38 -33
flask_app.py
CHANGED
|
@@ -6,6 +6,12 @@ import re
|
|
| 6 |
import tempfile
|
| 7 |
from flask import Flask, request, jsonify, render_template
|
| 8 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 9 |
# Disable CUDA and excessive parallel threads to save memory
|
| 10 |
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
| 11 |
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
@@ -44,70 +50,71 @@ def upload_file():
|
|
| 44 |
"""Route handling document upload, splitting, chunking, and vectorization."""
|
| 45 |
|
| 46 |
global retriever, LLM_model, api_key
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 55 |
api_key = request.form.get("apiKey")
|
| 56 |
if not api_key:
|
| 57 |
-
return "API key missing!", 400
|
| 58 |
|
| 59 |
uploaded = request.files.get("file")
|
| 60 |
if not uploaded or uploaded.filename.strip() == "":
|
| 61 |
-
return "No file uploaded", 400
|
| 62 |
|
| 63 |
ext = uploaded.filename.rsplit(".", 1)[-1].lower()
|
| 64 |
-
with tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}") as
|
| 65 |
-
uploaded.save(
|
| 66 |
-
path =
|
| 67 |
|
| 68 |
-
#
|
| 69 |
try:
|
| 70 |
loader = PyPDFLoader(path) if ext == "pdf" else TextLoader(path)
|
| 71 |
documents = loader.load()
|
| 72 |
except Exception as e:
|
| 73 |
os.unlink(path)
|
| 74 |
-
return f"Failed to read document: {e}", 400
|
| 75 |
|
| 76 |
if not documents:
|
| 77 |
os.unlink(path)
|
| 78 |
-
return "No readable content found in the document.", 400
|
| 79 |
-
|
| 80 |
|
| 81 |
-
#
|
| 82 |
-
splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
|
| 83 |
chunks = splitter.split_documents(documents)
|
| 84 |
|
| 85 |
-
|
| 86 |
-
# Light embedding model (fast + low memory)
|
| 87 |
try:
|
| 88 |
# embeds = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-MiniLM-L3-v2")
|
| 89 |
-
embeds = HuggingFaceEmbeddings(model_name="./models/paraphrase-MiniLM-L3-v2")
|
| 90 |
vector_store = FAISS.from_documents(chunks, embeds)
|
| 91 |
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
|
| 92 |
|
| 93 |
except Exception as e:
|
| 94 |
os.unlink(path)
|
| 95 |
-
return f"Embedding model failed: {e}", 500
|
| 96 |
-
|
| 97 |
|
| 98 |
-
# Initialize
|
| 99 |
try:
|
| 100 |
LLM_model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=api_key)
|
| 101 |
except Exception as e:
|
| 102 |
-
return f"Failed to initialize chat model: {e}", 500
|
| 103 |
-
|
| 104 |
|
| 105 |
-
# Cleanup
|
| 106 |
os.unlink(path)
|
| 107 |
del documents, chunks, vector_store
|
| 108 |
gc.collect()
|
| 109 |
|
| 110 |
-
return "Document processed successfully! You can now ask questions."
|
| 111 |
|
| 112 |
|
| 113 |
@app.route("/chat", methods=["POST"])
|
|
@@ -161,14 +168,12 @@ def chat():
|
|
| 161 |
response = f"Error generating response: {str(e)}"
|
| 162 |
|
| 163 |
# Clean markdown artifacts
|
| 164 |
-
cleaned = re.sub(r"
|
| 165 |
-
cleaned = re.sub(r"\*(.*?)\*", r"\1", cleaned)
|
| 166 |
|
| 167 |
gc.collect()
|
| 168 |
return jsonify({"answer": cleaned})
|
| 169 |
|
| 170 |
|
| 171 |
-
# run app
|
| 172 |
if __name__ == "__main__":
|
| 173 |
port = int(os.environ.get("PORT", 7860))
|
| 174 |
-
app.run(host="0.0.0.0", port=port, debug=False)
|
|
|
|
| 6 |
import tempfile
|
| 7 |
from flask import Flask, request, jsonify, render_template
|
| 8 |
|
| 9 |
+
# # Pre-download and save the embedding model
|
| 10 |
+
# from sentence_transformers import SentenceTransformer
|
| 11 |
+
# model = SentenceTransformer("sentence-transformers/paraphrase-MiniLM-L3-v2")
|
| 12 |
+
# model.save("models/paraphrase-MiniLM-L3-v2")
|
| 13 |
+
|
| 14 |
+
|
| 15 |
# Disable CUDA and excessive parallel threads to save memory
|
| 16 |
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
|
| 17 |
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
| 50 |
"""Route handling document upload, splitting, chunking, and vectorization."""
|
| 51 |
|
| 52 |
global retriever, LLM_model, api_key
|
| 53 |
+
|
| 54 |
+
try:
|
| 55 |
+
# Import heavy dependencies only when needed
|
| 56 |
+
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
| 57 |
+
from langchain_community.vectorstores import FAISS
|
| 58 |
+
from langchain_community.document_loaders import TextLoader, PyPDFLoader
|
| 59 |
+
from langchain_huggingface import HuggingFaceEmbeddings
|
| 60 |
+
from langchain_google_genai import ChatGoogleGenerativeAI
|
| 61 |
+
except Exception as e:
|
| 62 |
+
return jsonify({"error": f"Missing dependency: {e}"}), 500
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
# Get user API key
|
| 66 |
api_key = request.form.get("apiKey")
|
| 67 |
if not api_key:
|
| 68 |
+
return jsonify({"error": "API key missing!"}), 400
|
| 69 |
|
| 70 |
uploaded = request.files.get("file")
|
| 71 |
if not uploaded or uploaded.filename.strip() == "":
|
| 72 |
+
return jsonify({"error": "No file uploaded."}), 400
|
| 73 |
|
| 74 |
ext = uploaded.filename.rsplit(".", 1)[-1].lower()
|
| 75 |
+
with tempfile.NamedTemporaryFile(delete=False, suffix=f".{ext}") as tmp:
|
| 76 |
+
uploaded.save(tmp.name)
|
| 77 |
+
path = tmp.name
|
| 78 |
|
| 79 |
+
# Load document
|
| 80 |
try:
|
| 81 |
loader = PyPDFLoader(path) if ext == "pdf" else TextLoader(path)
|
| 82 |
documents = loader.load()
|
| 83 |
except Exception as e:
|
| 84 |
os.unlink(path)
|
| 85 |
+
return jsonify({"error": f"Failed to read document: {e}"}), 400
|
| 86 |
|
| 87 |
if not documents:
|
| 88 |
os.unlink(path)
|
| 89 |
+
return jsonify({"error": "No readable content found in the document."}), 400
|
|
|
|
| 90 |
|
| 91 |
+
# Split document into smaller chunks
|
| 92 |
+
splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
|
| 93 |
chunks = splitter.split_documents(documents)
|
| 94 |
|
| 95 |
+
# Create embeddings & vector store
|
|
|
|
| 96 |
try:
|
| 97 |
# embeds = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-MiniLM-L3-v2")
|
| 98 |
+
embeds = HuggingFaceEmbeddings(model_name="./models/paraphrase-MiniLM-L3-v2") # local model (offline)
|
| 99 |
vector_store = FAISS.from_documents(chunks, embeds)
|
| 100 |
retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
|
| 101 |
|
| 102 |
except Exception as e:
|
| 103 |
os.unlink(path)
|
| 104 |
+
return jsonify({"error": f"Embedding model failed: {e}"}), 500
|
|
|
|
| 105 |
|
| 106 |
+
# Initialize Gemini model
|
| 107 |
try:
|
| 108 |
LLM_model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=api_key)
|
| 109 |
except Exception as e:
|
| 110 |
+
return jsonify({"error": f"Failed to initialize chat model: {e}"}), 500
|
|
|
|
| 111 |
|
| 112 |
+
# Cleanup
|
| 113 |
os.unlink(path)
|
| 114 |
del documents, chunks, vector_store
|
| 115 |
gc.collect()
|
| 116 |
|
| 117 |
+
return jsonify({"message": "Document processed successfully! You can now ask questions."})
|
| 118 |
|
| 119 |
|
| 120 |
@app.route("/chat", methods=["POST"])
|
|
|
|
| 168 |
response = f"Error generating response: {str(e)}"
|
| 169 |
|
| 170 |
# Clean markdown artifacts
|
| 171 |
+
cleaned = re.sub(r"[*_`#]+", "", response)
|
|
|
|
| 172 |
|
| 173 |
gc.collect()
|
| 174 |
return jsonify({"answer": cleaned})
|
| 175 |
|
| 176 |
|
|
|
|
| 177 |
if __name__ == "__main__":
|
| 178 |
port = int(os.environ.get("PORT", 7860))
|
| 179 |
+
app.run(host="0.0.0.0", port=port, debug=False)
|