| import os |
| import uuid |
| import requests |
| import gradio as gr |
|
|
| |
| MODELS = { |
| "Dolphin Llama3 8B": "cognitivecomputations/dolphin-2.9-llama3-8b", |
| "Dolphin Mixtral 8x7B": "cognitivecomputations/dolphin-mixtral-8x7b", |
| "Dolphin Llama3 8B 256K": "cognitivecomputations/dolphin-llama3-8b-256k" |
| } |
|
|
| HF_API_TOKEN = os.environ.get("HF_TOKEN", "") |
|
|
| client_memory = {} |
| client_tokens = {} |
|
|
| def ensure_client(cid): |
| if cid not in client_memory: |
| client_memory[cid] = [] |
| client_tokens[cid] = 0 |
|
|
| def query_hf(model, prompt): |
| url = f"https://api-inference.huggingface.co/models/{model}" |
| headers = {"Authorization": f"Bearer {HF_API_TOKEN}"} |
| data = {"inputs": prompt, "parameters": {"max_new_tokens": 512}} |
|
|
| response = requests.post(url, headers=headers, json=data) |
| try: |
| return response.json()[0]["generated_text"] |
| except: |
| return f"Error: {response.text}" |
|
|
| def chat_fn(message, client_id, model_name): |
| if not client_id or client_id.strip() == "": |
| client_id = str(uuid.uuid4()) |
|
|
| ensure_client(client_id) |
|
|
| history = "" |
| for msg in client_memory[client_id]: |
| history += f"{msg['role']}: {msg['content']}\n" |
|
|
| prompt = history + f"user: {message}\nassistant:" |
|
|
| model = MODELS[model_name] |
| output = query_hf(model, prompt) |
|
|
| client_memory[client_id].append({"role": "user", "content": message}) |
| client_memory[client_id].append({"role": "assistant", "content": output}) |
|
|
| client_tokens[client_id] += len(output.split()) |
|
|
| return output, client_id, client_tokens[client_id] |
|
|
| with gr.Blocks() as app: |
| gr.Markdown("# 🔥 Dolphin Chat — HF API Powered") |
|
|
| with gr.Row(): |
| client_id = gr.Textbox(label="Client ID") |
| model_selector = gr.Dropdown( |
| choices=list(MODELS.keys()), |
| value="Dolphin Llama3 8B", |
| label="Model" |
| ) |
|
|
| chatbox = gr.Chatbot() |
| message_box = gr.Textbox(label="Message") |
| send_btn = gr.Button("Send") |
| tokens_used = gr.Number(label="Tokens Used") |
|
|
| send_btn.click( |
| chat_fn, |
| inputs=[message_box, client_id, model_selector], |
| outputs=[chatbox, client_id, tokens_used] |
| ) |
|
|
| app.launch() |
