app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llmlingua import PromptCompressor

app = FastAPI()

# Initialize LLMLingua
compressor = PromptCompressor(model_name="openai-community/gpt2", device_map="cpu")

class TextInput(BaseModel):
    text: str

@app.post("/compress")
async def compress_text(input: TextInput):
    try:
        compressed_text = compressor.compress_prompt(input.text)
        # Return raw data for now
        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"}