commit

main.py

@@ -16,15 +16,15 @@ from pprint import pprint
 import asyncio
 import importlib.util
 import traceback
-
 import sys
 import json
 import jsonschema
-# import aiosqlite
 from utils import extract_code
 import numpy as np
 import os
 import requests
+import secrets
+import urllib.parse
 
 app = FastAPI()
 
@@ -32,54 +32,42 @@ client_id = os.getenv("OAUTH_CLIENT_ID")
 client_secret = os.getenv("OAUTH_CLIENT_SECRET")
 space_host = os.getenv("SPACE_HOST")
 
-# DATABASE_FILE = "samples.db"
-
-
 client = OpenAI(
-  base_url="https://openrouter.ai/api/v1",
-  api_key=os.environ.get('OPENROUTER_KEY')
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.environ.get('OPENROUTER_KEY')
 )
 
-
-# async def setup_database():
-#     async with aiosqlite.connect(DATABASE_FILE) as db:
-#         await db.execute("""
-#             CREATE TABLE IF NOT EXISTS samples (
-#                 hash TEXT PRIMARY KEY,
-#                 data TEXT NOT NULL,
-#                 dataset TEXT NOT NULL
-#             )
-#         """)
-#         await db.commit()
-
-# async def insert_sample(hash: str, data: str, dataset: str):
-#     async with aiosqlite.connect(DATABASE_FILE) as db:
-#         # Check if a record with the same hash already exists
-#         cursor = await db.execute("SELECT COUNT(*) FROM samples WHERE hash = ?", (hash,))
-#         count = await cursor.fetchone()
-
-#         if count[0] == 0:
-#             # Insert the new record since it doesn't exist
-#             await db.execute("INSERT INTO samples (hash, data, dataset) VALUES (?, ?, ?)", (hash, data, dataset))
-#             await db.commit()
-#         else:
-#             # A record with the same hash already exists
-#             print("Record with the same hash already exists in the database.")
-
-# async def get_sample_by_hash(hash: str):
-#     async with aiosqlite.connect(DATABASE_FILE) as db:
-#         cursor = await db.execute("SELECT data, dataset FROM samples WHERE hash = ?", (hash,))
-#         row = await cursor.fetchone()
-#         return row
+state_queue_map = {}
 
 def is_sharegpt(sample):
-    schema={'$schema': 'http://json-schema.org/schema#', 'type': 'object', 'properties': {'conversations': {'type': 'array', 'items': {'type': 'object', 'properties': {'from': { 'type': 'string', 'enum': ['human', 'gpt', 'system'] }, 'value': {'type': 'string'}}, 'required': ['from', 'value']}}}, 'required': ['conversations']}
+    schema = {
+        '$schema': 'http://json-schema.org/schema#',
+        'type': 'object',
+        'properties': {
+            'conversations': {
+                'type': 'array',
+                'items': {
+                    'type': 'object',
+                    'properties': {
+                        'from': {
+                            'type': 'string',
+                            'enum': [
+                                'human',
+                                'gpt',
+                                'system']},
+                        'value': {
+                            'type': 'string'}},
+                    'required': [
+                        'from',
+                        'value']}}},
+        'required': ['conversations']}
     try:
         jsonschema.validate(instance=sample, schema=schema)
         return True
     except jsonschema.exceptions.ValidationError as e:
         return False
 
+
 def sha256(string):
     # Create a hashlib object for SHA-256
     sha256_hash = hashlib.sha256()
@@ -88,6 +76,7 @@ def sha256(string):
 
     return sha256_hash.hexdigest()
 
+
 def get_adapter_name(sample):
     builder = SchemaBuilder()
     builder.add_object(sample)
@@ -95,6 +84,7 @@ def get_adapter_name(sample):
 
     return sha256(json.dumps(schema))
 
+
 def has_adapter(sample):
     adapter_name = get_adapter_name(sample)
 
@@ -106,37 +96,26 @@ def has_adapter(sample):
 
     return True
 
+
 def auto_tranform(sample):
     adapter_name = get_adapter_name(sample)
     if not has_adapter(sample):
         create_adapter(sample, adapter_name)
 
     module_name = f"dataset_adapters.{adapter_name}"
-    spec = importlib.util.spec_from_file_location(module_name, f"dataset_adapters/{adapter_name}.py")
+    spec = importlib.util.spec_from_file_location(
+        module_name, f"dataset_adapters/{adapter_name}.py")
     dynamic_module = importlib.util.module_from_spec(spec)
     sys.modules[module_name] = dynamic_module
     spec.loader.exec_module(dynamic_module)
 
-    # Use the function from the dynamically imported module
     transformed_data = dynamic_module.transform_data(sample)
 
     if isinstance(transformed_data, list):
-        return {'conversations' : transformed_data}
-
+        return {'conversations': transformed_data}
 
     return transformed_data
 
-
-
-
-# def create_adapter(sample, adapter_name):
-#     builder = SchemaBuilder()
-#     builder.add_object(sample)
-#     schema = builder.to_schema()
-
-#     code_string = """def transform_data(data):
-#     raise Exception('')"""
-
     with open(f"dataset_adapters/{adapter_name}.py", 'w') as file:
         file.write(code_string)
 
@@ -146,7 +125,7 @@ def create_adapter(sample, adapter_name):
     builder.add_object(sample)
     schema = builder.to_schema()
 
-    prompt = f"""Make me minimal and efficient python code to convert data in the shape of 
+    prompt = f"""Make me minimal and efficient python code to convert data in the shape of
 
 initial data shape
 ==========➡️📑📐==========
@@ -177,65 +156,26 @@ For transforming the data you shall use python. Make robust and elegant python c
 
 your code will contain a function `def transform_data(data):` that does the transformation and outputs the newly shaped data. Only the data, no schema. Your code snippet will include only the function signature and body. I know how to call it. You won't need to import anything, I will take care of parsing and dumping json. You work with dicts. Remember to be careful if you iterate over the data because I want the output conversation to always start with the prompt. In other words, always process "input" before "output" and "instruction" before "output". Such heuristics are very important. If there is "instruction" and "input" and the "input" is not empty, concat the input at the end of the first message. If the data contains no "system" message, human always speaks first. If it contains a "system"  message, the "system" message is first, then human, then gpt, then alternating if needed
 
-"human" ALWAYS SPEAKS BEFORE "gpt", if you suspect your code makes "gpt speak first, fix it 
+"human" ALWAYS SPEAKS BEFORE "gpt", if you suspect your code makes "gpt speak first, fix it
 
 MOST IMPORTANT IS THAT YOU look at the initial data shape (➡️📑📐) to ground your transformation into final data shape (⬇️📑📐)
 
 Your output should contain only the code for `def transform_data(data):`, signature and body. Put the code inside markdown code block"""
 
     response = client.chat.completions.create(
-                        model="openai/gpt-4-1106-preview", # Optional (user controls the default)
-                        messages=[
-                        { "role": "system", "content": """You are ChatGPT, a large language model trained by OpenAI, based on the GPT-4 architecture.
+        model="openai/gpt-4-1106-preview",
+        messages=[
+            {"role": "system", "content": """You are ChatGPT, a large language model trained by OpenAI, based on the GPT-4 architecture.
 Knowledge cutoff: 2023-04
 Current date: 2023-11-05
 
-Image input capabilities: Enabled""" },
-#                             {"role": "user", "content": f"""Make me minimal and efficient python code to convert data in the shape of 
-
-# ```jsonschema
-# {json.dumps(schema)}
-# ```
-
-# to equivalent data in the form ```
-# {{'$schema': 'http://json-schema.org/schema#', 'type': 'object', 'properties': {{'conversations': {{'type': 'array', 'items': {{'type': 'object', 'properties': {{'from': {{ 'type': 'string', 'enum': ['human', 'gpt', 'system'] }}, 'value': {{'type': 'string'}}}}, 'required': ['from', 'value']}}}}}}, 'required': ['conversations']}}
-# ```
-
-# the input is 
-# ```
-# {json.dumps(sample)}
-# ```
-
-
-# `input` is usually associated with `"from" : "human"` while `output` is usually associated with `"from" : "gpt"`
-
-# don't transform, make robust and elegant python code that will do the transformation
-
-
-# your code will contain a function `def transform_data(data):` that does the transformation and outputs the newly shaped data. Only the data, no schema. Your code snippet will include only the function signature and body. I know how to call it. You won't need to import anything, I will take care of parsing and dumping json. You work with dicts. Remember to be careful if you iterate over the data because I want the output conversation to always start with the prompt. In other words, always process "input" before "output" and "instruction" before "output". Such heuristics are very important. If there is "instruction" and "input" and the "input" is not empty, concat the input at the end of the first message."""
-# }
-  {"role": "user", "content": prompt}
-                        ]
-                               )
+Image input capabilities: Enabled"""},
+            {"role": "user", "content": prompt}
+        ]
+    )
 
     val = response.choices[0].message.content
-    # index = val.index('def transform_data(data)')
-
-    # def get_code_start():
-    #     for i in range(index,0,-1):
-    #         if val[i:i+3] == "```":
-    #             idx = val[i:].index('\n')
-    #             return i + (idx) + 1
-
-    # def get_code_end():
-    #     for i in range(index, len(val)):
-    #         if val[i:i+3] == "```":
-    #             return i-1
 
-    # code_string = val[get_code_start():get_code_end()]
-
-
-    # print("###", val)
     code_string = extract_code(val)
 
     if code_string is None:
@@ -251,28 +191,59 @@ async def get_sample(hash: str = Query(..., alias="hash")):
     if res is None:
         raise HTTPException(status_code=404, detail="Item not found")
     data, dataset = res
-    sample= auto_tranform(json.loads(data))
+    sample = auto_tranform(json.loads(data))
     return {'sample': sample, 'dataset': dataset}
 
-@app.get("/random-sample-stream")
-async def get_random_sample(background_tasks: BackgroundTasks, dataset_name: str = Query(..., alias="dataset-name"), index: str = Query(None, alias="index")):
+
+def generate_random_string(length=16):
+    return secrets.token_hex(length)
+
+
+
+@app.get("/oauth_token")
+async def get_oauth_token():
     queue = asyncio.Queue()
-    def event_stream(queue):
-        yield f"data: {json.dumps({'status': 'grab_sample'})}\n\n" 
+
+    async def event_stream(queue, state):
+        state_queue_map[state] = queue
+
+        redirect_uri = f'https://{space_host}/login/callback'
+
+        auth_url = (
+            f"https://huggingface.co/oauth/authorize?"
+            f"redirect_uri={urllib.parse.quote(redirect_uri)}&"
+            f"client_id={client_id}&"
+            f"scope=openid%20profile&"
+            f"response_type=code&"
+            f"state={state}"
+        )
+        yield f"data: {json.dumps({ 'url' : auth_url })}\n\n"
+
         try:
+            while True:
+                message = await queue.get()
+                if 'end_stream' in message and message['end_stream']:
+                    break
+                yield f"data: {json.dumps(message)}\n\n"
+        finally:
+            del state_queue_map[state]
 
+    state = generate_random_string()
+    return StreamingResponse(
+        event_stream(queue, state),
+        media_type="text/event-stream")
 
 
+@app.get("/random-sample-stream")
+async def get_random_sample(background_tasks: BackgroundTasks, dataset_name: str = Query(..., alias="dataset-name"), index: str = Query(None, alias="index")):
+    queue = asyncio.Queue()
 
-            # dataset = load_dataset(dataset_name,streaming=True)
-            # split = [key for key in dataset.keys() if "train" in key]
-            
-            
-            
-            
-            import requests
+    def event_stream(queue):
+        yield f"data: {json.dumps({'status': 'grab_sample'})}\n\n"
+        try:
             headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN')}"}
             API_URL = f"https://datasets-server.huggingface.co/info?dataset={dataset_name}"
+
             def query():
                 response = requests.get(API_URL, headers=headers)
                 return response.json()
@@ -283,15 +254,15 @@ async def get_random_sample(background_tasks: BackgroundTasks, dataset_name: str
 
             num_samples = split['num_examples']
             split_name = split['name']
-            
-            # dataset = load_dataset(dataset_name, split=split_name, streaming=True)
-            idx = random.randint(0, num_samples) if index is None else int(index)
 
+            idx = random.randint(
+                0, num_samples) if index is None else int(index)
 
             API_URL = f"https://datasets-server.huggingface.co/rows?dataset={dataset_name}&config=default&split=train&offset={idx}&length=1"
-          
+
             def query():
-                headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN')}"}
+                headers = {
+                    "Authorization": f"Bearer {os.environ.get('HF_TOKEN')}"}
                 response = requests.get(API_URL, headers=headers)
 
                 if response.status_code != 200:
@@ -301,15 +272,7 @@ async def get_random_sample(background_tasks: BackgroundTasks, dataset_name: str
 
             random_sample = data['rows'][0]['row']
 
-            # pprint(random_sample)
-
-
-            # selected = dataset.skip(idx)
-            # random_sample = next(iter(selected))#random.choice(samples_buffer)
-
             hashed = sha256(json.dumps(random_sample))
-            # insert_sample(hashed, json.dumps(random_sample), dataset_name)
-            # background_tasks.add_task(insert_sample, hashed, json.dumps(random_sample), dataset_name)
 
         except Exception as e:
             message = ""
@@ -317,9 +280,9 @@ async def get_random_sample(background_tasks: BackgroundTasks, dataset_name: str
                 message = e.message
             else:
                 message = str(e)
-                
+
             print("error : ", message)
-            yield f"data: {json.dumps({'status': 'error', 'message' : message })}\n\n" 
+            yield f"data: {json.dumps({'status': 'error', 'message' : message })}\n\n"
 
         transformed_data = random_sample
 
@@ -328,7 +291,7 @@ async def get_random_sample(background_tasks: BackgroundTasks, dataset_name: str
         if not is_sharegpt(random_sample):
             try:
                 if not has_adapter(random_sample):
-                    yield f"data: {json.dumps({'status': 'creating_adapter'})}\n\n" 
+                    yield f"data: {json.dumps({'status': 'creating_adapter'})}\n\n"
 
                 transformed_data = auto_tranform(random_sample)
             except Exception as e:
@@ -337,27 +300,29 @@ async def get_random_sample(background_tasks: BackgroundTasks, dataset_name: str
                     print("error : ", e.message)
                 else:
                     print("error : ", e)
-                yield f"data: {json.dumps({'status': 'error'})}\n\n" 
+                yield f"data: {json.dumps({'status': 'error'})}\n\n"
 
         if success:
-            yield f"data: {json.dumps({'status': 'done', 'data' : transformed_data, 'index' : str(idx)})}\n\n" 
-
-    return StreamingResponse(event_stream(queue), media_type="text/event-stream")
+            yield f"data: {json.dumps({'status': 'done', 'data' : transformed_data, 'index' : str(idx)})}\n\n"
 
+    return StreamingResponse(
+        event_stream(queue),
+        media_type="text/event-stream")
 
 
 @app.get("/random-sample")
 async def get_random_sample(dataset_name: str = Query(..., alias="dataset-name")):
     try:
-        dataset = load_dataset(dataset_name,streaming=True)
+        dataset = load_dataset(dataset_name, streaming=True)
         split = [key for key in dataset.keys() if "train" in key]
         dataset = load_dataset(dataset_name, split=split[0], streaming=True)
 
         buffer_size = 100  # Define a reasonable buffer size
-        samples_buffer = [sample for _, sample in zip(range(buffer_size), dataset)]
-        
-        random_sample = random.choice(samples_buffer)
+        samples_buffer = [
+            sample for _, sample in zip(
+                range(buffer_size), dataset)]
 
+        random_sample = random.choice(samples_buffer)
 
         hashed = sha256(json.dumps(random_sample))
 
@@ -368,12 +333,12 @@ async def get_random_sample(dataset_name: str = Query(..., alias="dataset-name")
         if module_spec is None:
             create_adapter(random_sample, sanitized)
 
-        spec = importlib.util.spec_from_file_location(module_name, f"dataset_adapters/{sanitized}.py")
+        spec = importlib.util.spec_from_file_location(
+            module_name, f"dataset_adapters/{sanitized}.py")
         dynamic_module = importlib.util.module_from_spec(spec)
         sys.modules[module_name] = dynamic_module
         spec.loader.exec_module(dynamic_module)
 
-        # Use the function from the dynamically imported module
         transformed_data = dynamic_module.transform_data(random_sample)
 
         return transformed_data
@@ -384,16 +349,14 @@ async def get_random_sample(dataset_name: str = Query(..., alias="dataset-name")
         raise HTTPException(status_code=500, detail=str(e))
 
 
-
 @app.get("/login/callback")
 async def oauth_callback(code: str, state: str):
-    # Prepare the authorization header
     credentials = f"{client_id}:{client_secret}"
     credentials_bytes = credentials.encode("ascii")
     base64_credentials = base64.b64encode(credentials_bytes)
     auth_header = f"Basic {base64_credentials.decode('ascii')}"
     username = ""
-    
+
     try:
         token_response = requests.post(
             'https://huggingface.co/oauth/token',
@@ -411,31 +374,26 @@ async def oauth_callback(code: str, state: str):
             tokens = token_response.json()
             access_token = tokens.get('access_token')
 
-
             if access_token:
-                    try:
-                        user_response = requests.get(
-                            'https://huggingface.co/api/user',
-                            headers={'Authorization': f'Bearer {access_token}'}
-                        )
-
-                        if user_response.status_code == 200:
-                            user_data = user_response.json()
-                            username = user_data['username']
-                            # Now you have the username, you can handle it as needed
-                        else:
-                            username = ""
-                            print(f"Error getting user data: {user_response.status_code}, {user_response.text}")
-
-                    except Exception:
-                        traceback.print_exc()
-                        username = ""
-            else:
-                username = ""
+                url = "https://huggingface.co/oauth/userinfo"
+
+                payload = ""
+                headers = {
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {access_token}"
+                }
+
+                response = requests.request(
+                    "GET", url, data=payload, headers=headers)
 
+                print(response.text)
 
-            # ID Token can be extracted here if needed
-            # id_token = tokens.get('id_token')
+                if state in state_queue_map:
+                    queue = state_queue_map[state]
+                    await queue.put({"access_token": access_token, "username": username})
+                    await queue.put({"end_stream": True})
+            else:
+                username = ""
         else:
             access_token = ""
 
@@ -443,23 +401,20 @@ async def oauth_callback(code: str, state: str):
         traceback.print_exc()
         access_token = ""
 
-    return {"access_token": access_token, "username" : username}
+    return {"access_token": access_token, "username": username}
+
 
 @app.get("/oauth-config")
 async def get_oauth_config(request: Request):
-    # client_host = "https://huggingface.co/spaces/thomasgauthier/ChatExplorer#request.client.host
     return {
         "client_id": client_id,
         "redirect_uri": f'https://{space_host}/login/callback'
     }
 
 
-# # @app.on_event("startup")
-# # async def startup_event():
-# #     await setup_database()
 @app.get("/")
 def index() -> FileResponse:
     return FileResponse(path="static/index.html", media_type="text/html")
 
 
-app.mount("/", StaticFiles(directory="static"), name="static")
+app.mount("/", StaticFiles(directory="static"), name="static")

static/index.html

@@ -5,7 +5,7 @@
     <link rel="icon" type="image/svg+xml" href="/vite.svg" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Vite + Preact</title>
-    <script type="module" crossorigin src="/assets/index-66194b32.js"></script>
+    <script type="module" crossorigin src="/assets/index-7974ca0c.js"></script>
     <link rel="stylesheet" href="/assets/index-abe6d7fb.css">
   </head>
   <body ondrop="event.preventDefault()" >