import gradio as gr
from docarray import Document
from jina import Flow
import clip_server
from clip_client import Client
import asyncio

def setupServer():
  cas_path = clip_server.__path__[0]
  flow_yaml = f'''
  jtype: Flow
  with:
    protocol: grpc
    port: 51000
  executors:
    - name: clip_t
      uses:
        jtype: CLIPEncoder
        metas:
          py_modules:
            - {cas_path}/executors/clip_torch.py
  '''
  
  f = Flow.load_config(flow_yaml)
  f.start()

def setupClient():
  c = Client('https://demo-cas.jina.ai:8443')
  return c

# setupServer()
c = setupClient()

def rank(uri, options):
   r = c.rank([
        Document(
            uri=uri,
            matches=[
                Document(text=f'a photo of a {p}')
                for p in options.split(',')
            ],
        )
    ])
   return r['@m', ['text', 'scores__clip_score__value']]

examples = [["https://picsum.photos/id/102/300/300", "three berry, four berries, ten berries"]]
iface = gr.Interface(fn=rank, inputs=["text", "text"], outputs="text", examples=examples)
iface.launch()