app.py

import gradio as gr
import pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM
import os

# Hugging Face API 토큰을 환경 변수로 설정 (로그인 후 토큰을 얻은 후 사용)
os.environ["HF_TOKEN"] = "your_huggingface_token"


# Parquet 파일 로드
data_file_path = "train-00000-of-00001-7f15f39e4c3a7ee9.parquet"  # 파일이 동일한 디렉토리에 있다고 가정
dataset = pd.read_parquet(data_file_path)

# 모델과 토크나이저 로드
model_id = "CohereForAI/c4ai-command-r-plus-08-2024"
tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=os.getenv("HF_TOKEN"))
model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=os.getenv("HF_TOKEN"))

# 질문-답변 데이터셋 (필드 이름이 'Question', 'Answer'라고 가정)
questions = dataset['Question'].tolist()
answers = dataset['Answer'].tolist()

# LLM을 사용하여 질문에 대한 답변 생성
def get_answer(user_question):
    # 입력 메시지를 채팅 템플릿에 맞춰 포맷팅
    messages = [{"role": "user", "content": user_question}]
    input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
    
    # 답변 생성
    gen_tokens = model.generate(
        input_ids, 
        max_new_tokens=100, 
        do_sample=True, 
        temperature=0.3,
    )
    
    # 생성된 텍스트 디코딩
    gen_text = tokenizer.decode(gen_tokens[0])
    return gen_text

# Gradio UI 구성
iface = gr.Interface(
    fn=get_answer,  # 실행할 함수
    inputs="text",  # 입력은 텍스트
    outputs="text",  # 출력은 텍스트
    title="의료 상담 AI",
    description="질문을 입력하고 실행 버튼을 누르면 LLM을 사용해 답변을 생성합니다.",
)

# UI 실행
iface.launch()