import gradio as gr
import requests
import os

# 환경 변수에서 Hugging Face API 토큰을 가져옵니다.
ACCESS_TOKEN = os.getenv("HF_TOKEN")

def respond(message, system_message, max_tokens=512, temperature=0.7, top_p=0.95):
    # 메시지를 구성
    data = {
        "model": "command-r-plus:104b-fp16",
        "prompt": message,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p
    }

    # 외부 API로 요청을 보냄
    response = requests.post("http://hugpu.ai:7877/api/generate", json=data)
    generated_text = response.json().get('generated_text', '')  # API 응답에서 텍스트 추출

    return generated_text  # 생성된 텍스트 반환


demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-P",
        ),
        
    ],
    fill_height=True,
    chatbot=chatbot,
    theme="Nymbo/Nymbo_Theme",
)
if __name__ == "__main__":
    demo.launch()