llm_lingua / app.py
Oluwatoni's picture
Update app.py
5ca4c2b verified
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llmlingua import PromptCompressor
app = FastAPI()
# Initialize LLMLingua-2 for better compression
compressor = PromptCompressor(model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",use_llmlingua2=True
, device_map="cpu")
class TextInput(BaseModel):
text: str
@app.post("/compress")
async def compress_text(input: TextInput):
try:
# Compress the input text using LLMLingua-2
compressed_text = compressor.compress_prompt(input.text)
# Ensure the compressed text is of high quality
if not compressed_text:
raise HTTPException(status_code=400, detail="Compression failed, no output generated.")
return {"compressed_text": compressed_text}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
async def root():
return {"message": "Welcome to the LLMLingua Text Compression API"}