app.py

#load package
from fastapi import FastAPI,HTTPException
from pydantic import BaseModel
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    StoppingCriteria,
    StoppingCriteriaList,
    TextIteratorStreamer
)
from typing import List, Tuple
from threading import Thread
import os
from pydantic import BaseModel
import logging
import uvicorn


# Configurer les répertoires de cache
os.environ['TRANSFORMERS_CACHE'] = '/app/.cache'
os.environ['HF_HOME'] = '/app/.cache'
# Charger le modèle et le tokenizer
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)


#Additional information
 
Informations = """ 
-text : Texte à resumé 

output:
- Text summary : texte resumé
"""

app =FastAPI(
    title='Text Summary',
    description =Informations
)  
 
#class to define the input text 
logging.basicConfig(level=logging.INFO)
logger =logging.getLogger(__name__)

@app.get("/")
async def home():
    return 'STN BIG DATA'

# 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")

# 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 :

1. Le problème principal
2. Les détails pertinents
3. L'impact sur le plaignant
4. Toute action ou résolution demandée

Résumez la plainte suivante en 3-4 phrases :

"""

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)