from fastapi import FastAPI,HTTPException,status,UploadFile,File
from pydantic import BaseModel
import uvicorn
import logging
import torch
import ollama
import os
import numpy as np

os.environ['TRANSFORMERS_CACHE'] = '/app/.cache'
os.environ['HF_HOME'] = '/app/.cache'
 
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__)


prompt_first ="Résumé en 4 ou 5 phrases de la plainte suivante :"
class TextSummary(BaseModel): 
    prompt:str 
    
@app.get("/")
async def home():
    return 'STN BIG DATA'
# Fonction pour générer du texte à partir d'une requête
# Modèle pour la requête
class RequestModel(BaseModel):
    text: str
@app.post("/generate/")
async def generate_text(request: RequestModel):
    # Assurer la configuration pour une utilisation optimale de la mémoire
    try:
        response = ollama.chat(
            model="llama3:latest",
            messages=[{"role": "user", "content": prompt_first + request.text}],
        )
        
        generated_text = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                generated_text += chunk.choices[0].delta.content

        # Traitement du texte pour supprimer les phrases génériques
        intro_phrases = [
            "Voici un résumé de la plainte en 5 phrases :",
            "Résumé :",
            "Voici ce qui s'est passé :",
            "Cette plainte a été déposée par"
        ]
        
        for phrase in intro_phrases:
            if generated_text.startswith(phrase):
                generated_text = generated_text[len(phrase):].strip()
                break

        return {"generated_text": generated_text}
    except ollama._types.ResponseError as e:
        return {"error": str(e)}
if __name__ == "__main__":
    uvicorn.run("app:app",reload=True)