File size: 5,836 Bytes
018c8f2 0619e92 018c8f2 0619e92 018c8f2 0619e92 546aaea 0619e92 018c8f2 0619e92 0fc7fa2 018c8f2 0619e92 018c8f2 0619e92 018c8f2 0619e92 0fc7fa2 018c8f2 0619e92 018c8f2 0619e92 018c8f2 0619e92 018c8f2 0619e92 018c8f2 0619e92 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | # Fichier: app.py (VERSION CORRIGÉE FINALE - OPTIMISÉE POUR LA MÉMOIRE DU SPACE)
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch
import os
# --- Configuration du Modèle ---
#model_id = "microsoft/Phi-3-mini-4k-instruct"
model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
device = torch.device("cpu")
# --- Stratégie de chargement pour économiser la mémoire (Quantisation) ---
# Si le Space a un GPU/CUDA, la quantisation sera utilisée, réduisant la RAM par 8.
# Si le Space est CPU seulement, cette tentative échouera, et nous utiliserons le fallback float32.
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
# Charger le Tokenizer et le Modèle
try:
tokenizer = AutoTokenizer.from_pretrained(model_id)
# TENTATIVE 1 : Chargement avec Quantisation 4-bit (Méthode recommandée)
print("Tentative de chargement avec quantisation 4-bit...")
# Le chargement en 4-bit nécessite device_map="auto"
model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=quantization_config,
device_map="auto",
trust_remote_code=True
)
print(f"Modèle {model_id} chargé et quantifié.")
except Exception as e_quant:
# Si la quantisation échoue (souvent sans GPU), on revient à la version CPU
print(f"Échec de la quantisation : {e_quant}. Tentative de chargement float32 CPU (Attention: peut causer OOM).")
# TENTATIVE 2 : Fallback sur le chargement float32 CPU (Votre code initial, mais avec fix du bug)
try:
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float32,
trust_remote_code=True
).to(device)
print(f"Modèle {model_id} chargé sur CPU (Float32).")
except Exception as e_cpu:
print(f"Échec critique du chargement CPU : {e_cpu}")
# Si même float32 échoue, vous avez BESOIN de plus de RAM pour votre Space.
raise e_cpu
model.eval()
app = FastAPI(
title="NLP Space - Phi-3 Mini API (CPU)",
description="API REST pour génération, résumé et classification de texte, optimisée pour CPU."
)
# Schéma de données pour les requêtes POST
class PromptRequest(BaseModel):
prompt: str
max_tokens: int = 500
temperature: float = 0.7
# Fonction utilitaire pour interagir avec le modèle
def generate_text_from_model(system_prompt: str, user_prompt: str, max_tokens: int, temperature: float):
# Formatage de l'instruction pour Phi-3 Instruct (Chat Template)
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
text_to_generate = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
# Trouver le device réel du modèle pour y placer les inputs (nécessaire après le chargement device_map)
# Assurez-vous que le modèle est correctement placé, en le forçant sur CPU si nécessaire.
real_device = model.device if model.device.type != 'meta' else torch.device("cpu")
inputs = tokenizer(text_to_generate, return_tensors="pt").to(real_device)
with torch.no_grad():
output = model.generate(
**inputs,
max_new_tokens=max_tokens,
do_sample=True,
temperature=temperature,
pad_token_id=tokenizer.eos_token_id,
use_cache=False # CORRECTION CRITIQUE 2: Fixe le bug DynamicCache
)
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
# Nettoyage
response_start_tag = "<|assistant|>"
if response_start_tag in generated_text:
return generated_text.split(response_start_tag, 1)[1].strip()
return generated_text.strip()
# --- Endpoints (Inchangés) ---
@app.post("/generate")
async def generate(request: PromptRequest):
"""Génération de texte libre."""
system_prompt = "Tu es un assistant IA très utile et créatif."
try:
result = generate_text_from_model(
system_prompt=system_prompt,
user_prompt=request.prompt,
max_tokens=request.max_tokens,
temperature=request.temperature
)
return {"result": result}
except Exception as e:
return {"error": str(e)}
@app.post("/summarize")
async def summarize(request: PromptRequest):
# ... (code inchangé) ...
system_prompt = "Tu es un expert en résumé concis et précis. Ton objectif est de résumer le texte fourni de manière à en conserver l'idée principale."
user_prompt = f"Résume le texte suivant de manière concise et factuelle:\n\n---\n\n{request.prompt}"
try:
result = generate_text_from_model(
system_prompt=system_prompt,
user_prompt=user_prompt,
max_tokens=request.max_tokens,
temperature=0.3
)
return {"result": result}
except Exception as e:
return {"error": str(e)}
@app.post("/classify")
async def classify(request: PromptRequest):
# ... (code inchangé) ...
system_prompt = "Tu es un expert en classification. Réponds uniquement avec l'étiquette de classification demandée sans phrases supplémentaires."
user_prompt = request.prompt
try:
result = generate_text_from_model(
system_prompt=system_prompt,
user_prompt=user_prompt,
max_tokens=50,
temperature=0.1
)
return {"result": result}
except Exception as e:
return {"error": str(e)} |