fixe

app.py

@@ -1,5 +1,5 @@
 #load package
-from fastapi import FastAPI
+from fastapi import FastAPI,HTTPException
 from pydantic import BaseModel
 import torch
 from transformers import (
@@ -43,88 +43,51 @@ app =FastAPI(
 logging.basicConfig(level=logging.INFO)
 logger =logging.getLogger(__name__)
 
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_ids = model.config.eos_token_id
-        for stop_id in stop_ids:
-            if input_ids[0][-1] == stop_id:
-                return True
-        return False
-
-default_prompt = """Bonjour,
-
-En tant qu’expert en gestion des plaintes réseaux, rédige un descriptif clair de la plainte ci-dessous. Résume la situation en 4 ou 5 phrases concises, en mettant l'accent sur :
-1. **Informations Client** : Indique des détails pertinents sur le client.
-2. **Dates et Délais** : Mentionne les dates clés et les délais (prise en charge, résolution, etc.).
-3. **Contexte et Détails** : Inclut les éléments essentiels de la plainte (titre, détails, états d’avancement, qualification, fichiers joints).
-
-Ajoute une recommandation importante pour éviter le mécontentement du client, par exemple, en cas de service non fourni malgré le paiement. Adapte le ton pour qu'il soit humain et engageant.
-
-Merci !
-
-"""
-class PredictionRequest(BaseModel):
-    history: list
-    prompt: str
-    max_length: int = 128000
-    top_p: float = 0.8
-    temperature: float = 0.6
-
-class PredictionResponse(BaseModel):
-    history: list
-
 @app.get("/")
 async def home():
     return 'STN BIG DATA'
-@app.post("/predict", response_model=PredictionResponse)
-async def predict(request: PredictionRequest):
-    stop = StopOnTokens()
-    messages = []
-    query = ""
-
-    # Préparer les messages sans inclure le prompt par défaut dans l'historique renvoyé
-    for idx, (user_msg, model_msg) in enumerate(request.history):
-        if idx == 0 and request.prompt:  # Ignorer le prompt par défaut dans l'historique
-            continue
-        if idx == len(request.history) - 1 and not model_msg:
-            query = user_msg
-            break
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if model_msg:
-            messages.append({"role": "assistant", "content": model_msg})
 
-    # Inclure le prompt uniquement pour la génération, mais pas dans l'historique
-    model_inputs = tokenizer.build_chat_input(query, history=messages, role='user').input_ids.to(
-        next(model.parameters()).device
-    )
+# Charger le modèle et le tokenizer
+model_name = "THUDM/longwriter-glm4-9b"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="auto")
 
-    eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
-                    tokenizer.get_command("<|observation|>")]
+# Prompt par défaut
+default_prompt = """Vous êtes un assistant expert en résumé de plaintes. Votre tâche est de résumer la plainte fournie de manière concise et professionnelle, en incluant les points clés suivants :
 
-    generate_kwargs = {
-        "input_ids": model_inputs,
-        "max_new_tokens": request.max_length,
-        "do_sample": True,
-        "top_p": request.top_p,
-        "temperature": request.temperature,
-        "stopping_criteria": StoppingCriteriaList([stop]),
-        "repetition_penalty": 1,
-        "eos_token_id": eos_token_id,
-    }
+1. Le problème principal
+2. Les détails pertinents
+3. L'impact sur le plaignant
+4. Toute action ou résolution demandée
 
-    # Générer le texte
-    output = model.generate(**generate_kwargs)
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+Résumez la plainte suivante en 3-4 phrases :
 
-    # Ajouter le texte généré à l'historique
-    request.history[-1][1] += generated_text
+"""
 
-    # Retourner l'historique sans le prompt
-    return PredictionResponse(history=request.history)
+class ComplaintInput(BaseModel):
+    text: str
+
+@app.post("/summarize_complaint")
+async def summarize_complaint(input: ComplaintInput):
+    try:
+        full_prompt = default_prompt + input.text
+        inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=150,
+                num_return_sequences=1,
+                no_repeat_ngram_size=2,
+                temperature=0.7
+            )
+        
+        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Enlever le prompt initial de la sortie
+        summary = summary.replace(full_prompt, "").strip()
+        return {"summary": summary}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    uvicorn.run("app:app",reload=True) 
-
-
- 
+    uvicorn.run("app:app",reload=True)