|
|
|
from fastapi import FastAPI |
|
from pydantic import BaseModel |
|
import torch |
|
from transformers import ( |
|
AutoModelForCausalLM, |
|
AutoTokenizer |
|
) |
|
from typing import List, Tuple |
|
from threading import Thread |
|
import os |
|
from pydantic import BaseModel |
|
import logging |
|
import uvicorn |
|
|
|
|
|
|
|
os.environ['TRANSFORMERS_CACHE'] = '/app/.cache' |
|
os.environ['HF_HOME'] = '/app/.cache' |
|
|
|
model = AutoModelForCausalLM.from_pretrained("THUDM/longwriter-glm4-9b", trust_remote_code=True, device_map='auto') |
|
tokenizer = AutoTokenizer.from_pretrained("THUDM/longwriter-glm4-9b", trust_remote_code=True) |
|
|
|
|
|
|
|
|
|
Informations = """ |
|
-text : Texte à resumé |
|
|
|
output: |
|
- Text summary : texte resumé |
|
""" |
|
|
|
app =FastAPI( |
|
title='Text Summary', |
|
description =Informations |
|
) |
|
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 ! |
|
|
|
""" |
|
|
|
logging.basicConfig(level=logging.INFO) |
|
logger =logging.getLogger(__name__) |
|
|
|
class PredictionRequest(BaseModel): |
|
text: str = None |
|
max_length: int = 512 |
|
|
|
@app.post("/predict/") |
|
async def predict(request: PredictionRequest): |
|
|
|
if request.text: |
|
prompt = default_prompt + "\n\n" + request.text |
|
else: |
|
prompt = default_prompt |
|
|
|
|
|
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device) |
|
outputs = model.generate(inputs, max_length=request.max_length, do_sample=True) |
|
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
return {"generated_text": generated_text} |
|
|
|
if __name__ == "__main__": |
|
uvicorn.run("app:app",reload=True) |
|
|
|
|
|
|