# app.py
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread
import gradio as gr
import re
import torch
from openai import OpenAI

client = OpenAI(
    api_key="sk-420ab66020704eabbe37501ec39b7a2b",
    base_url="https://bailingchat.alipay.com",
)

# define chat function
def chat(user_input, max_tokens=11264):
    # chat history
    messages_template = [
  #      {"role": "system", "content": "You are Ling, an assistant created by inclusionAI"},
        {"role": "user", "content": user_input}
    ]

    response = client.chat.completions.create(
        model="Ling-lite-1.5-250604",
        messages=messages_template,
        max_tokens=max_tokens,
        temperature=0.01,
        top_p=1,
    )
    resp_text = response.choices[0].message.content
    print(resp_text)
    yield resp_text



# Create a custom layout using Blocks
with gr.Blocks(css="""
    #markdown-output {
        height: 300px;
        overflow-y: auto;
        border: 1px solid #ddd;
        padding: 10px;
    }
""") as demo:
    gr.Markdown(
        "## Ling-lite-1.5 AI Assistant\n"
        "Based on [inclusionAI/Ling-lite-1.5](https://huggingface.co/inclusionAI/Ling-lite-1.5) "
    )

    with gr.Row():
        max_tokens_slider = gr.Slider(minimum=100, maximum=10000, step=100, label="Generated length")

#    output_box = gr.Textbox(lines=10, label="Response")
    output_box = gr.Markdown(label="Response", elem_id="markdown-output")
    input_box = gr.Textbox(lines=8, label="Input you question")

    examples = gr.Examples(
        examples=[
            ["Introducing the basic concepts of large language models"],
            ["How to solve long context dependencies in math problems?"]
        ],
        inputs=input_box
    )

    interface = gr.Interface(
        fn=chat,
        inputs=[input_box, max_tokens_slider],
        outputs=output_box,
        live=False  # disable auto-triggering on input change
    )

# launch Gradio Service
demo.queue()
demo.launch()

# Construct Gradio Interface
#interface = gr.Interface(
#    fn=chat,
#    inputs=[
#        gr.Textbox(lines=8, label="输入你的问题"),
#        gr.Slider(minimum=100, maximum=102400, step=50, label="生成长度")
#    ],
#    outputs=[
#        gr.Textbox(lines=8, label="模型回复")
#    ],
#    title="Ling-lite-1.5 AI助手",
#    description="基于 [inclusionAI/Ling-lite-1.5](https://huggingface.co/inclusionAI/Ling-lite-1.5)  的对话式文本生成演示。",
#    examples=[
#        ["介绍大型语言模型的基本概念"],
#        ["如何解决数学问题中的长上下文依赖？"]
#    ]
#)

# launch Gradion Service
#interface.launch()