KJ24 commited on
Commit
5f710b6
·
verified ·
1 Parent(s): 42cc3fb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -24
app.py CHANGED
@@ -4,44 +4,43 @@ from fastapi import FastAPI
4
  from pydantic import BaseModel
5
  from typing import Optional
6
 
7
- # ✅ Modules LlamaIndex (version >= 0.10.0)
8
  from llama_index.core import Document
9
  from llama_index.core.settings import Settings
10
- from llama_index.core.node_parser import SemanticSplitterNodeParser
11
- from llama_index.core.text_splitter import RecursiveTextSplitter
12
  from llama_index.llms.llama_cpp import LlamaCPP
13
  from llama_index.core.base.llms.base import BaseLLM
14
 
15
- # ✅ Embedding local (transformers + torch)
16
  from transformers import AutoTokenizer, AutoModel
17
  import torch
18
  import torch.nn.functional as F
19
  import os
20
 
21
- # ✅ Initialisation de l'app FastAPI
22
  app = FastAPI()
23
 
24
- # ✅ Configuration du cache Hugging Face (important pour HF Spaces)
25
  CACHE_DIR = "/app/cache"
26
  os.environ["HF_HOME"] = CACHE_DIR
27
  os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
28
  os.environ["HF_MODULES_CACHE"] = CACHE_DIR
29
  os.environ["HF_HUB_CACHE"] = CACHE_DIR
30
 
31
- # ✅ Choix du modèle d'embedding dense (ex : BGE-small)
32
  MODEL_NAME = "BAAI/bge-small-en-v1.5"
33
  tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
34
  model = AutoModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
35
 
36
- # ✅ Fonction d'embedding normalisé (vectorisation dense)
37
  def get_embedding(text: str):
 
38
  with torch.no_grad():
39
  inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
40
  outputs = model(**inputs)
41
- embeddings = outputs.last_hidden_state[:, 0] # On prend le token [CLS]
42
  return F.normalize(embeddings, p=2, dim=1).squeeze().tolist()
43
 
44
- # ✅ Format des données envoyées à l’API
45
  class ChunkRequest(BaseModel):
46
  text: str
47
  max_tokens: Optional[int] = 1000
@@ -51,13 +50,12 @@ class ChunkRequest(BaseModel):
51
  source: Optional[str] = None
52
  type: Optional[str] = None
53
 
54
- # ✅ Route de l’API pour le chunking sémantique
55
  @app.post("/chunk")
56
  async def chunk_text(data: ChunkRequest):
57
  try:
58
  print(f"\n✅ Texte reçu ({len(data.text)} caractères) : {data.text[:200]}...", flush=True)
59
 
60
- # ✅ Chargement du modèle GGUF distant avec LlamaCPP (CPU friendly)
61
  llm = LlamaCPP(
62
  model_url="https://huggingface.co/TheBloke/CodeLlama-7B-Instruct-GGUF/resolve/main/codellama-7b-instruct.Q4_K_M.gguf",
63
  temperature=0.1,
@@ -67,35 +65,42 @@ async def chunk_text(data: ChunkRequest):
67
  model_kwargs={"n_gpu_layers": 1},
68
  )
69
 
70
- print("✅ Modèle LLM chargé avec succès !")
71
 
72
- # ✅ Wrapper embedding compatible avec LlamaIndex
73
  class SimpleEmbedding:
74
  def get_text_embedding(self, text: str):
75
  return get_embedding(text)
76
 
77
- # ✅ Configuration globale de LlamaIndex
78
- assert isinstance(llm, BaseLLM), "❌ L’objet LLM nest pas compatible avec LlamaIndex"
79
  Settings.llm = llm
80
  Settings.embed_model = SimpleEmbedding()
81
 
82
- print("✅ Configuration du LLM et de l'embedding terminée. On initialise le Semantic Splitter...")
83
 
84
- parser = SemanticSplitterNodeParser.from_defaults(llm=llm)
85
  doc = Document(text=data.text)
86
 
 
 
 
87
  try:
88
  nodes = parser.get_nodes_from_documents([doc])
89
  print(f"✅ Semantic Splitter : {len(nodes)} chunks générés")
90
  if not nodes:
91
- raise ValueError("Aucun chunk produit par SemanticSplitter")
 
92
  except Exception as e:
93
- print(f"⚠️ Fallback vers RecursiveTextSplitter suite à : {e}")
94
- splitter = RecursiveTextSplitter(chunk_size=data.max_tokens, chunk_overlap=data.overlap)
 
 
 
95
  nodes = splitter.get_nodes_from_documents([doc])
96
  print(f"♻️ Recursive Splitter : {len(nodes)} chunks générés")
97
 
98
- # ✅ Résultat structuré pour n8n ou autre client HTTP
99
  return {
100
  "chunks": [node.text for node in nodes],
101
  "metadatas": [node.metadata for node in nodes],
@@ -103,14 +108,14 @@ async def chunk_text(data: ChunkRequest):
103
  "titre": data.titre,
104
  "source": data.source,
105
  "type": data.type,
106
- "error": None # utilisé par n8n pour signaler "pas d'erreur"
107
  }
108
 
109
  except Exception as e:
110
  print(f"❌ Erreur critique : {e}")
111
  return {"error": str(e)}
112
 
113
- # ✅ Lancement local (facultatif pour HF Spaces)
114
  if __name__ == "__main__":
115
  import uvicorn
116
  uvicorn.run("app:app", host="0.0.0.0", port=7860)
 
4
  from pydantic import BaseModel
5
  from typing import Optional
6
 
7
+ # ✅ LlamaIndex (version >= 0.10.0)
8
  from llama_index.core import Document
9
  from llama_index.core.settings import Settings
10
+ from llama_index.core.node_parser import SemanticSplitterNodeParser, RecursiveCharacterTextSplitter
 
11
  from llama_index.llms.llama_cpp import LlamaCPP
12
  from llama_index.core.base.llms.base import BaseLLM
13
 
14
+ # ✅ Embedding local (basé sur transformers + torch)
15
  from transformers import AutoTokenizer, AutoModel
16
  import torch
17
  import torch.nn.functional as F
18
  import os
19
 
20
+ # ✅ Initialisation de l'application FastAPI
21
  app = FastAPI()
22
 
23
+ # ✅ Configuration du cache local de Hugging Face pour économiser l'espace dans le container
24
  CACHE_DIR = "/app/cache"
25
  os.environ["HF_HOME"] = CACHE_DIR
26
  os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
27
  os.environ["HF_MODULES_CACHE"] = CACHE_DIR
28
  os.environ["HF_HUB_CACHE"] = CACHE_DIR
29
 
30
+ # ✅ Modèle d'embedding local utilisé pour vectoriser les textes
31
  MODEL_NAME = "BAAI/bge-small-en-v1.5"
32
  tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
33
  model = AutoModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
34
 
 
35
  def get_embedding(text: str):
36
+ """Fonction pour générer un embedding dense normalisé à partir d’un texte."""
37
  with torch.no_grad():
38
  inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
39
  outputs = model(**inputs)
40
+ embeddings = outputs.last_hidden_state[:, 0] # On prend le vecteur [CLS]
41
  return F.normalize(embeddings, p=2, dim=1).squeeze().tolist()
42
 
43
+ # ✅ Schéma des données attendues dans le POST
44
  class ChunkRequest(BaseModel):
45
  text: str
46
  max_tokens: Optional[int] = 1000
 
50
  source: Optional[str] = None
51
  type: Optional[str] = None
52
 
 
53
  @app.post("/chunk")
54
  async def chunk_text(data: ChunkRequest):
55
  try:
56
  print(f"\n✅ Texte reçu ({len(data.text)} caractères) : {data.text[:200]}...", flush=True)
57
 
58
+ # ✅ Chargement du modèle LLM CodeLlama quantifié (GGUF) via URL Hugging Face
59
  llm = LlamaCPP(
60
  model_url="https://huggingface.co/TheBloke/CodeLlama-7B-Instruct-GGUF/resolve/main/codellama-7b-instruct.Q4_K_M.gguf",
61
  temperature=0.1,
 
65
  model_kwargs={"n_gpu_layers": 1},
66
  )
67
 
68
+ print("✅ Modèle CodeLlama-7B chargé avec succès !")
69
 
70
+ # ✅ Embedding local pour LlamaIndex
71
  class SimpleEmbedding:
72
  def get_text_embedding(self, text: str):
73
  return get_embedding(text)
74
 
75
+ # ✅ Configuration du moteur dans LlamaIndex
76
+ assert isinstance(llm, BaseLLM), "❌ Le LLM n'est pas compatible avec Settings.llm"
77
  Settings.llm = llm
78
  Settings.embed_model = SimpleEmbedding()
79
 
80
+ print("✅ Configuration du LLM et de l'embedding terminée.")
81
 
82
+ # Document à découper
83
  doc = Document(text=data.text)
84
 
85
+ # ✅ Split intelligent (semantic)
86
+ parser = SemanticSplitterNodeParser.from_defaults(llm=llm)
87
+
88
  try:
89
  nodes = parser.get_nodes_from_documents([doc])
90
  print(f"✅ Semantic Splitter : {len(nodes)} chunks générés")
91
  if not nodes:
92
+ raise ValueError("Aucun chunk généré par SemanticSplitter")
93
+
94
  except Exception as e:
95
+ print(f"⚠️ Fallback vers RecursiveCharacterTextSplitter suite à : {e}")
96
+ splitter = RecursiveCharacterTextSplitter(
97
+ chunk_size=data.max_tokens,
98
+ chunk_overlap=data.overlap
99
+ )
100
  nodes = splitter.get_nodes_from_documents([doc])
101
  print(f"♻️ Recursive Splitter : {len(nodes)} chunks générés")
102
 
103
+ # ✅ Construction de la réponse
104
  return {
105
  "chunks": [node.text for node in nodes],
106
  "metadatas": [node.metadata for node in nodes],
 
108
  "titre": data.titre,
109
  "source": data.source,
110
  "type": data.type,
111
+ "error": None # utile pour n8n ou tout autre client
112
  }
113
 
114
  except Exception as e:
115
  print(f"❌ Erreur critique : {e}")
116
  return {"error": str(e)}
117
 
118
+ # ✅ Lancement du serveur si exécution directe (mode debug)
119
  if __name__ == "__main__":
120
  import uvicorn
121
  uvicorn.run("app:app", host="0.0.0.0", port=7860)