File size: 794 Bytes
67b6614
 
 
 
 
 
 
8836a04
 
67b6614
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8836a04
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llmlingua import PromptCompressor
import os

app = FastAPI()

# Initialize LLM Lingua with CPU
compressor = PromptCompressor( device_map="auto")  # or set device_map=None

class TextInput(BaseModel):
    text: str

class CompressedOutput(BaseModel):
    compressed_text: str

@app.post("/compress", response_model=CompressedOutput)
async def compress_text(input: TextInput):
    try:
        compressed = compressor.compress_prompt(input.text)
        return CompressedOutput(compressed_text=compressed)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))