KJ24 commited on
Commit
ec7f6a1
·
verified ·
1 Parent(s): d42f1d2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -34
app.py CHANGED
@@ -2,15 +2,40 @@ from fastapi import FastAPI
2
  from pydantic import BaseModel
3
  from typing import Optional
4
 
 
5
  from llama_index.core.settings import Settings
6
  from llama_index.core import Document
7
- from llama_index.embeddings.huggingface import HuggingFaceEmbedding
8
  from llama_index.llms.llama_cpp import LlamaCPP
9
  from llama_index.core.node_parser import SemanticSplitterNodeParser
10
 
 
 
 
 
 
 
11
  app = FastAPI()
12
 
13
- # 📥 Modèle de la requête JSON envoyée à /chunk
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
  class ChunkRequest(BaseModel):
15
  text: str
16
  source_id: Optional[str] = None
@@ -20,38 +45,26 @@ class ChunkRequest(BaseModel):
20
 
21
  @app.post("/chunk")
22
  async def chunk_text(data: ChunkRequest):
23
- # ✅ Chargement direct d’un modèle hébergé sur Hugging Face (pas de fichier local .gguf)
24
- llm = LlamaCPP(
25
- model_url="https://huggingface.co/leafspark/Mistral-7B-Instruct-v0.2-Q4_K_M-GGUF/resolve/main/mistral-7b-instruct-v0.2.Q4_K_M.gguf",
26
- temperature=0.1,
27
- max_new_tokens=512,
28
- context_window=2048,
29
- generate_kwargs={"top_p": 0.95},
30
- model_kwargs={"n_gpu_layers": 1}, # Laisse 1 si CPU
31
- )
32
-
33
- # ✅ Embedding open-source via Hugging Face
34
- embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
35
-
36
- # ✅ Configuration du service IA
37
- # service_context = ServiceContext.from_defaults(
38
- # llm=llm,
39
- # embed_model=embed_model
40
- # )
41
-
42
 
43
- # ✅ Nouvelle méthode recommandée : paramétrer Settings globalement
44
- Settings.llm = llm
45
- Settings.embed_model = embed_model
 
46
 
 
 
47
 
48
-
49
- try:
50
  # ✅ Découpage sémantique intelligent
51
- # parser = SemanticSplitterNodeParser.from_defaults(service_context=service_context)
52
-
53
- # ✅ Appel du parser sans service_context
54
-
55
  parser = SemanticSplitterNodeParser.from_defaults()
56
  nodes = parser.get_nodes_from_documents([Document(text=data.text)])
57
 
@@ -61,14 +74,12 @@ async def chunk_text(data: ChunkRequest):
61
  "source_id": data.source_id,
62
  "titre": data.titre,
63
  "source": data.source,
64
- "type": data.type
65
  }
 
66
  except Exception as e:
67
  return {"error": str(e)}
68
 
69
-
70
  if __name__ == "__main__":
71
  import uvicorn
72
  uvicorn.run("app:app", host="0.0.0.0", port=7860)
73
-
74
-
 
2
  from pydantic import BaseModel
3
  from typing import Optional
4
 
5
+ # ✅ Modules de LlamaIndex
6
  from llama_index.core.settings import Settings
7
  from llama_index.core import Document
 
8
  from llama_index.llms.llama_cpp import LlamaCPP
9
  from llama_index.core.node_parser import SemanticSplitterNodeParser
10
 
11
+ # ✅ Pour l'embedding LOCAL via transformers
12
+ from transformers import AutoTokenizer, AutoModel
13
+ import torch
14
+ import torch.nn.functional as F
15
+ import os
16
+
17
  app = FastAPI()
18
 
19
+ # Configuration locale du cache HF pour Hugging Face
20
+ CACHE_DIR = "/data"
21
+ os.environ["HF_HOME"] = CACHE_DIR
22
+ os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
23
+ os.environ["HF_MODULES_CACHE"] = CACHE_DIR
24
+ os.environ["HF_HUB_CACHE"] = CACHE_DIR
25
+
26
+ # ✅ Configuration du modèle d’embedding local (ex: BGE / Nomic / GTE etc.)
27
+ MODEL_NAME = "BAAI/bge-small-en-v1.5"
28
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
29
+ model = AutoModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
30
+
31
+ def get_embedding(text: str):
32
+ with torch.no_grad():
33
+ inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
34
+ outputs = model(**inputs)
35
+ embeddings = outputs.last_hidden_state[:, 0]
36
+ return F.normalize(embeddings, p=2, dim=1).squeeze().tolist()
37
+
38
+ # ✅ Données entrantes du POST
39
  class ChunkRequest(BaseModel):
40
  text: str
41
  source_id: Optional[str] = None
 
45
 
46
  @app.post("/chunk")
47
  async def chunk_text(data: ChunkRequest):
48
+ try:
49
+ # Chargement du modèle LLM depuis Hugging Face en ligne (pas de .gguf local)
50
+ llm = LlamaCPP(
51
+ model_url="https://huggingface.co/leafspark/Mistral-7B-Instruct-v0.2-Q4_K_M-GGUF/resolve/main/mistral-7b-instruct-v0.2.Q4_K_M.gguf",
52
+ temperature=0.1,
53
+ max_new_tokens=512,
54
+ context_window=2048,
55
+ generate_kwargs={"top_p": 0.95},
56
+ model_kwargs={"n_gpu_layers": 1},
57
+ )
 
 
 
 
 
 
 
 
 
58
 
59
+ # ✅ Intégration manuelle de l'embedding local dans Settings
60
+ class SimpleEmbedding:
61
+ def get_text_embedding(self, text: str):
62
+ return get_embedding(text)
63
 
64
+ Settings.llm = llm
65
+ Settings.embed_model = SimpleEmbedding()
66
 
 
 
67
  # ✅ Découpage sémantique intelligent
 
 
 
 
68
  parser = SemanticSplitterNodeParser.from_defaults()
69
  nodes = parser.get_nodes_from_documents([Document(text=data.text)])
70
 
 
74
  "source_id": data.source_id,
75
  "titre": data.titre,
76
  "source": data.source,
77
+ "type": data.type,
78
  }
79
+
80
  except Exception as e:
81
  return {"error": str(e)}
82
 
 
83
  if __name__ == "__main__":
84
  import uvicorn
85
  uvicorn.run("app:app", host="0.0.0.0", port=7860)