from fastapi import FastAPI
from fastapi import File, UploadFile
import shutil
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from os import makedirs,getcwd
from os.path import join,exists,dirname
from modelscope.models import Model
from modelscope.pipelines import pipeline

model = Model.from_pretrained('damo/multi-modal_convnext-roberta-base_vldoc-embedding')
doc_VL_emb_pipeline = pipeline(task='document-vl-embedding', model=model)

app = FastAPI()
parent_path  = dirname(getcwd())

temp_path = join(parent_path,'temp')
if not exists(temp_path):
    makedirs(temp_path)

@app.post("/analyze")
def pdf2images(file: UploadFile=File(...)):
    file_savePath =  join(temp_path,file.filename)

    with open(file_savePath,'wb') as f:
        shutil.copyfileobj(file.file, f)

    inp = {
        'images': ['/demo.png'], 
        'ocr_info_paths': ['/demo.json']
    }
    result = doc_VL_emb_pipeline(inp)
    
    print('Results of VLDoc: ')
    for k, v in result.items():
        print(f'{k}: {v.size()}')
    return result