Spaces:
Running
on
Zero
Running
on
Zero
import torch | |
import gradio as gr | |
import spaces | |
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
import os | |
from threading import Thread | |
import random | |
from datasets import load_dataset | |
import numpy as np | |
from sklearn.feature_extraction.text import TfidfVectorizer | |
import pandas as pd | |
from typing import List, Tuple | |
import json | |
from datetime import datetime | |
import pyarrow.parquet as pq | |
import pypdf | |
import io | |
import pyarrow.parquet as pq | |
from pdfminer.high_level import extract_text | |
from pdfminer.layout import LAParams | |
from tabulate import tabulate # tabulate ์ถ๊ฐ | |
import platform | |
import subprocess | |
import pytesseract | |
from pdf2image import convert_from_path | |
# ์ ์ญ ๋ณ์ ์ถ๊ฐ | |
current_file_context = None | |
# ํ๊ฒฝ ๋ณ์ ์ค์ | |
HF_TOKEN = os.environ.get("HF_TOKEN", None) | |
MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024" | |
MODELS = os.environ.get("MODELS") | |
MODEL_NAME = MODEL_ID.split("/")[-1] | |
model = None # ์ ์ญ ๋ณ์๋ก ์ ์ธ | |
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) | |
# ์ํคํผ๋์ ๋ฐ์ดํฐ์ ๋ก๋ | |
wiki_dataset = load_dataset("lcw99/wikipedia-korean-20240501-1million-qna") | |
print("Wikipedia dataset loaded:", wiki_dataset) | |
# TF-IDF ๋ฒกํฐ๋ผ์ด์ ์ด๊ธฐํ ๋ฐ ํ์ต | |
print("TF-IDF ๋ฒกํฐํ ์์...") | |
questions = wiki_dataset['train']['question'][:10000] # ์ฒ์ 10000๊ฐ๋ง ์ฌ์ฉ | |
vectorizer = TfidfVectorizer(max_features=1000) | |
question_vectors = vectorizer.fit_transform(questions) | |
print("TF-IDF ๋ฒกํฐํ ์๋ฃ") | |
class ChatHistory: | |
def __init__(self): | |
self.history = [] | |
self.history_file = "/tmp/chat_history.json" | |
self.load_history() | |
def add_conversation(self, user_msg: str, assistant_msg: str): | |
conversation = { | |
"timestamp": datetime.now().isoformat(), | |
"messages": [ | |
{"role": "user", "content": user_msg}, | |
{"role": "assistant", "content": assistant_msg} | |
] | |
} | |
self.history.append(conversation) | |
self.save_history() | |
def format_for_display(self): | |
formatted = [] | |
for conv in self.history: | |
formatted.append([ | |
conv["messages"][0]["content"], | |
conv["messages"][1]["content"] | |
]) | |
return formatted | |
def get_messages_for_api(self): | |
messages = [] | |
for conv in self.history: | |
messages.extend([ | |
{"role": "user", "content": conv["messages"][0]["content"]}, | |
{"role": "assistant", "content": conv["messages"][1]["content"]} | |
]) | |
return messages | |
def clear_history(self): | |
self.history = [] | |
self.save_history() | |
def save_history(self): | |
try: | |
with open(self.history_file, 'w', encoding='utf-8') as f: | |
json.dump(self.history, f, ensure_ascii=False, indent=2) | |
except Exception as e: | |
print(f"ํ์คํ ๋ฆฌ ์ ์ฅ ์คํจ: {e}") | |
def load_history(self): | |
try: | |
if os.path.exists(self.history_file): | |
with open(self.history_file, 'r', encoding='utf-8') as f: | |
self.history = json.load(f) | |
except Exception as e: | |
print(f"ํ์คํ ๋ฆฌ ๋ก๋ ์คํจ: {e}") | |
self.history = [] | |
# ์ ์ญ ChatHistory ์ธ์คํด์ค ์์ฑ | |
chat_history = ChatHistory() | |
def find_relevant_context(query, top_k=3): | |
# ์ฟผ๋ฆฌ ๋ฒกํฐํ | |
query_vector = vectorizer.transform([query]) | |
# ์ฝ์ฌ์ธ ์ ์ฌ๋ ๊ณ์ฐ | |
similarities = (query_vector * question_vectors.T).toarray()[0] | |
# ๊ฐ์ฅ ์ ์ฌํ ์ง๋ฌธ๋ค์ ์ธ๋ฑ์ค | |
top_indices = np.argsort(similarities)[-top_k:][::-1] | |
# ๊ด๋ จ ์ปจํ ์คํธ ์ถ์ถ | |
relevant_contexts = [] | |
for idx in top_indices: | |
if similarities[idx] > 0: | |
relevant_contexts.append({ | |
'question': questions[idx], | |
'answer': wiki_dataset['train']['answer'][idx], | |
'similarity': similarities[idx] | |
}) | |
return relevant_contexts | |
def analyze_file_content(content, file_type): | |
"""Analyze file content and return structural summary""" | |
if file_type in ['parquet', 'csv']: | |
try: | |
lines = content.split('\n') | |
header = lines[0] | |
columns = header.count('|') - 1 | |
rows = len(lines) - 3 | |
return f"๐ ๋ฐ์ดํฐ์ ๊ตฌ์กฐ: {columns}๊ฐ ์ปฌ๋ผ, {rows}๊ฐ ๋ฐ์ดํฐ" | |
except: | |
return "โ ๋ฐ์ดํฐ์ ๊ตฌ์กฐ ๋ถ์ ์คํจ" | |
lines = content.split('\n') | |
total_lines = len(lines) | |
non_empty_lines = len([line for line in lines if line.strip()]) | |
if any(keyword in content.lower() for keyword in ['def ', 'class ', 'import ', 'function']): | |
functions = len([line for line in lines if 'def ' in line]) | |
classes = len([line for line in lines if 'class ' in line]) | |
imports = len([line for line in lines if 'import ' in line or 'from ' in line]) | |
return f"๐ป ์ฝ๋ ๊ตฌ์กฐ: {total_lines}์ค (ํจ์: {functions}, ํด๋์ค: {classes}, ์ํฌํธ: {imports})" | |
paragraphs = content.count('\n\n') + 1 | |
words = len(content.split()) | |
return f"๐ ๋ฌธ์ ๊ตฌ์กฐ: {total_lines}์ค, {paragraphs}๋จ๋ฝ, ์ฝ {words}๋จ์ด" | |
def extract_pdf_text_with_ocr(file_path): | |
try: | |
# Poppler ๊ฒฝ๋ก ์ค์ | |
if platform.system() == 'Windows': | |
poppler_path = r"C:\Program Files\poppler-0.68.0\bin" | |
else: | |
poppler_path = None # Linux์ ๊ฒฝ์ฐ ๊ธฐ๋ณธ ๊ฒฝ๋ก ์ฌ์ฉ | |
# PDF๋ฅผ ์ด๋ฏธ์ง๋ก ๋ณํ | |
images = convert_from_path( | |
file_path, | |
poppler_path=poppler_path, | |
fmt='jpeg', | |
grayscale=False, | |
size=(1700, None) # ํด์๋ ํฅ์ | |
) | |
# ์ ์ฒด ํ ์คํธ ์ ์ฅ | |
text = "" | |
# ๊ฐ ํ์ด์ง์ ๋ํด OCR ์ํ | |
for i, image in enumerate(images): | |
try: | |
# OCR ์ค์ | |
custom_config = r'--oem 3 --psm 6 -l kor+eng' | |
# OCR ์ํ | |
page_text = pytesseract.image_to_string( | |
image, | |
config=custom_config | |
) | |
text += f"\n--- ํ์ด์ง {i+1} ---\n{page_text}\n" | |
except Exception as e: | |
print(f"ํ์ด์ง {i+1} OCR ์ค๋ฅ: {str(e)}") | |
continue | |
return text | |
except Exception as e: | |
return f"PDF ํ ์คํธ ์ถ์ถ ์ค๋ฅ: {str(e)}" | |
def read_uploaded_file(file): | |
if file is None: | |
return "", "" | |
try: | |
file_ext = os.path.splitext(file.name)[1].lower() | |
# Parquet ํ์ผ ์ฒ๋ฆฌ | |
if file_ext == '.parquet': | |
try: | |
table = pq.read_table(file.name) | |
df = table.to_pandas() | |
content = f"๐ Parquet ํ์ผ ๋ถ์:\n\n" | |
content += f"1. ๊ธฐ๋ณธ ์ ๋ณด:\n" | |
content += f"- ์ ์ฒด ํ ์: {len(df):,}๊ฐ\n" | |
content += f"- ์ ์ฒด ์ด ์: {len(df.columns)}๊ฐ\n" | |
content += f"- ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋: {df.memory_usage(deep=True).sum() / 1024 / 1024:.2f} MB\n\n" | |
content += f"2. ์ปฌ๋ผ ์ ๋ณด:\n" | |
for col in df.columns: | |
content += f"- {col} ({df[col].dtype})\n" | |
content += f"\n3. ๋ฐ์ดํฐ ๋ฏธ๋ฆฌ๋ณด๊ธฐ:\n" | |
# tabulate ์ฌ์ฉํ์ฌ ํ ์ด๋ธ ํ์์ผ๋ก ์ถ๋ ฅ | |
content += tabulate(df.head(5), headers='keys', tablefmt='pipe', showindex=False) | |
content += f"\n\n4. ๊ฒฐ์ธก์น ์ ๋ณด:\n" | |
null_counts = df.isnull().sum() | |
for col, count in null_counts[null_counts > 0].items(): | |
content += f"- {col}: {count:,}๊ฐ ({count/len(df)*100:.1f}%)\n" | |
# ์์นํ ์ปฌ๋ผ์ ๋ํ ๊ธฐ๋ณธ ํต๊ณ | |
numeric_cols = df.select_dtypes(include=['int64', 'float64']).columns | |
if len(numeric_cols) > 0: | |
content += f"\n5. ์์นํ ์ปฌ๋ผ ํต๊ณ:\n" | |
stats_df = df[numeric_cols].describe() | |
content += tabulate(stats_df, headers='keys', tablefmt='pipe') | |
return content, "parquet" | |
except Exception as e: | |
return f"Parquet ํ์ผ ์ฝ๊ธฐ ์ค๋ฅ: {str(e)}", "error" | |
# PDF ํ์ผ ์ฒ๋ฆฌ | |
if file_ext == '.pdf': | |
try: | |
pdf_reader = pypdf.PdfReader(file.name) | |
total_pages = len(pdf_reader.pages) | |
content = f"๐ PDF ๋ฌธ์ ๋ถ์:\n\n" | |
content += f"1. ๊ธฐ๋ณธ ์ ๋ณด:\n" | |
content += f"- ์ด ํ์ด์ง ์: {total_pages}ํ์ด์ง\n" | |
# ๋ฉํ๋ฐ์ดํฐ ์ถ์ถ | |
if pdf_reader.metadata: | |
content += "\n2. ๋ฉํ๋ฐ์ดํฐ:\n" | |
for key, value in pdf_reader.metadata.items(): | |
if value and str(key).startswith('/'): | |
content += f"- {key[1:]}: {value}\n" | |
# ๋จผ์ pdfminer๋ก ํ ์คํธ ์ถ์ถ ์๋ | |
try: | |
text = extract_text( | |
file.name, | |
laparams=LAParams( | |
line_margin=0.5, | |
word_margin=0.1, | |
char_margin=2.0, | |
all_texts=True | |
) | |
) | |
except: | |
text = "" | |
# pdfminer๋ก ์ถ์ถ ์คํจ์ OCR ์๋ | |
if not text.strip(): | |
text = extract_pdf_text_with_ocr(file.name) | |
# ํ ์คํธ ๋ถ์ | |
if text: | |
words = text.split() | |
lines = text.split('\n') | |
content += f"\n3. ํ ์คํธ ๋ถ์:\n" | |
content += f"- ์ด ๋จ์ด ์: {len(words):,}๊ฐ\n" | |
content += f"- ๊ณ ์ ๋จ์ด ์: {len(set(words)):,}๊ฐ\n" | |
content += f"- ์ด ๋ผ์ธ ์: {len(lines):,}๊ฐ\n" | |
# ๋ณธ๋ฌธ ๋ด์ฉ | |
content += f"\n4. ๋ณธ๋ฌธ ๋ด์ฉ:\n" | |
preview_length = min(2000, len(text)) # ๋ฏธ๋ฆฌ๋ณด๊ธฐ ๊ธธ์ด ์ฆ๊ฐ | |
content += f"--- ์ฒ์ {preview_length}์ ---\n" | |
content += text[:preview_length] | |
if len(text) > preview_length: | |
content += f"\n... (์ด {len(text):,}์ ์ค ์ผ๋ถ ํ์)\n" | |
else: | |
content += "\nโ ๏ธ ํ ์คํธ ์ถ์ถ ์คํจ" | |
return content, "pdf" | |
except Exception as e: | |
return f"PDF ํ์ผ ์ฝ๊ธฐ ์ค๋ฅ: {str(e)}", "error" | |
# CSV ํ์ผ ์ฒ๋ฆฌ | |
elif file_ext == '.csv': | |
encodings = ['utf-8', 'cp949', 'euc-kr', 'latin1'] | |
for encoding in encodings: | |
try: | |
df = pd.read_csv(file.name, encoding=encoding) | |
content = f"๐ CSV ํ์ผ ๋ถ์:\n\n" | |
content += f"1. ๊ธฐ๋ณธ ์ ๋ณด:\n" | |
content += f"- ์ ์ฒด ํ ์: {len(df):,}๊ฐ\n" | |
content += f"- ์ ์ฒด ์ด ์: {len(df.columns)}๊ฐ\n" | |
content += f"- ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋: {df.memory_usage(deep=True).sum() / 1024 / 1024:.2f} MB\n\n" | |
content += f"2. ์ปฌ๋ผ ์ ๋ณด:\n" | |
for col in df.columns: | |
content += f"- {col} ({df[col].dtype})\n" | |
content += f"\n3. ๋ฐ์ดํฐ ๋ฏธ๋ฆฌ๋ณด๊ธฐ:\n" | |
content += df.head(5).to_markdown(index=False) | |
content += f"\n\n4. ๊ฒฐ์ธก์น ์ ๋ณด:\n" | |
null_counts = df.isnull().sum() | |
for col, count in null_counts[null_counts > 0].items(): | |
content += f"- {col}: {count:,}๊ฐ ({count/len(df)*100:.1f}%)\n" | |
return content, "csv" | |
except UnicodeDecodeError: | |
continue | |
raise UnicodeDecodeError(f"์ง์๋๋ ์ธ์ฝ๋ฉ์ผ๋ก ํ์ผ์ ์ฝ์ ์ ์์ต๋๋ค ({', '.join(encodings)})") | |
# ํ ์คํธ ํ์ผ ์ฒ๋ฆฌ | |
else: | |
encodings = ['utf-8', 'cp949', 'euc-kr', 'latin1'] | |
for encoding in encodings: | |
try: | |
with open(file.name, 'r', encoding=encoding) as f: | |
content = f.read() | |
# ํ์ผ ๋ด์ฉ ๋ถ์ | |
lines = content.split('\n') | |
total_lines = len(lines) | |
non_empty_lines = len([line for line in lines if line.strip()]) | |
# ์ฝ๋ ํ์ผ ์ฌ๋ถ ํ์ธ | |
is_code = any(keyword in content.lower() for keyword in ['def ', 'class ', 'import ', 'function']) | |
analysis = f"\n๐ ํ์ผ ๋ถ์:\n" | |
if is_code: | |
# ์ฝ๋ ํ์ผ ๋ถ์ | |
functions = len([line for line in lines if 'def ' in line]) | |
classes = len([line for line in lines if 'class ' in line]) | |
imports = len([line for line in lines if 'import ' in line or 'from ' in line]) | |
analysis += f"- ํ์ผ ์ ํ: ์ฝ๋\n" | |
analysis += f"- ์ ์ฒด ๋ผ์ธ ์: {total_lines:,}์ค\n" | |
analysis += f"- ํจ์ ์: {functions}๊ฐ\n" | |
analysis += f"- ํด๋์ค ์: {classes}๊ฐ\n" | |
analysis += f"- import ๋ฌธ ์: {imports}๊ฐ\n" | |
else: | |
# ์ผ๋ฐ ํ ์คํธ ํ์ผ ๋ถ์ | |
words = len(content.split()) | |
chars = len(content) | |
analysis += f"- ํ์ผ ์ ํ: ํ ์คํธ\n" | |
analysis += f"- ์ ์ฒด ๋ผ์ธ ์: {total_lines:,}์ค\n" | |
analysis += f"- ์ค์ ๋ด์ฉ์ด ์๋ ๋ผ์ธ ์: {non_empty_lines:,}์ค\n" | |
analysis += f"- ๋จ์ด ์: {words:,}๊ฐ\n" | |
analysis += f"- ๋ฌธ์ ์: {chars:,}๊ฐ\n" | |
return content + analysis, "text" | |
except UnicodeDecodeError: | |
continue | |
raise UnicodeDecodeError(f"์ง์๋๋ ์ธ์ฝ๋ฉ์ผ๋ก ํ์ผ์ ์ฝ์ ์ ์์ต๋๋ค ({', '.join(encodings)})") | |
except Exception as e: | |
return f"ํ์ผ ์ฝ๊ธฐ ์ค๋ฅ: {str(e)}", "error" | |
# ํ์ผ ์ ๋ก๋ ์ด๋ฒคํธ ํธ๋ค๋ง ์์ | |
def init_msg(): | |
return "ํ์ผ์ ๋ถ์ํ๊ณ ์์ต๋๋ค..." | |
CSS = """ | |
/* 3D ์คํ์ผ CSS */ | |
:root { | |
--primary-color: #2196f3; | |
--secondary-color: #1976d2; | |
--background-color: #f0f2f5; | |
--card-background: #ffffff; | |
--text-color: #333333; | |
--shadow-color: rgba(0, 0, 0, 0.1); | |
} | |
body { | |
background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%); | |
min-height: 100vh; | |
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; | |
} | |
.container { | |
transform-style: preserve-3d; | |
perspective: 1000px; | |
} | |
.chatbot { | |
background: var(--card-background); | |
border-radius: 20px; | |
box-shadow: | |
0 10px 20px var(--shadow-color), | |
0 6px 6px var(--shadow-color); | |
transform: translateZ(0); | |
transition: transform 0.3s ease; | |
backdrop-filter: blur(10px); | |
} | |
.chatbot:hover { | |
transform: translateZ(10px); | |
} | |
/* ๋ฉ์์ง ์ ๋ ฅ ์์ญ */ | |
.input-area { | |
background: var(--card-background); | |
border-radius: 15px; | |
padding: 15px; | |
margin-top: 20px; | |
box-shadow: | |
0 5px 15px var(--shadow-color), | |
0 3px 3px var(--shadow-color); | |
transform: translateZ(0); | |
transition: all 0.3s ease; | |
display: flex; | |
align-items: center; | |
gap: 10px; | |
} | |
.input-area:hover { | |
transform: translateZ(5px); | |
} | |
/* ๋ฒํผ ์คํ์ผ */ | |
.custom-button { | |
background: linear-gradient(145deg, var(--primary-color), var(--secondary-color)); | |
color: white; | |
border: none; | |
border-radius: 10px; | |
padding: 10px 20px; | |
font-weight: 600; | |
cursor: pointer; | |
transform: translateZ(0); | |
transition: all 0.3s ease; | |
box-shadow: | |
0 4px 6px var(--shadow-color), | |
0 1px 3px var(--shadow-color); | |
} | |
.custom-button:hover { | |
transform: translateZ(5px) translateY(-2px); | |
box-shadow: | |
0 7px 14px var(--shadow-color), | |
0 3px 6px var(--shadow-color); | |
} | |
/* ํ์ผ ์ ๋ก๋ ๋ฒํผ */ | |
.file-upload-icon { | |
background: linear-gradient(145deg, #64b5f6, #42a5f5); | |
color: white; | |
border-radius: 8px; | |
font-size: 2em; | |
cursor: pointer; | |
display: flex; | |
align-items: center; | |
justify-content: center; | |
height: 70px; | |
width: 70px; | |
transition: all 0.3s ease; | |
box-shadow: 0 2px 5px rgba(0,0,0,0.1); | |
} | |
.file-upload-icon:hover { | |
transform: translateY(-2px); | |
box-shadow: 0 4px 8px rgba(0,0,0,0.2); | |
} | |
/* ํ์ผ ์ ๋ก๋ ๋ฒํผ ๋ด๋ถ ์์ ์คํ์ผ๋ง */ | |
.file-upload-icon > .wrap { | |
display: flex !important; | |
align-items: center; | |
justify-content: center; | |
width: 100%; | |
height: 100%; | |
} | |
.file-upload-icon > .wrap > p { | |
display: none !important; | |
} | |
.file-upload-icon > .wrap::before { | |
content: "๐"; | |
font-size: 2em; | |
display: block; | |
} | |
/* ๋ฉ์์ง ์คํ์ผ */ | |
.message { | |
background: var(--card-background); | |
border-radius: 15px; | |
padding: 15px; | |
margin: 10px 0; | |
box-shadow: | |
0 4px 6px var(--shadow-color), | |
0 1px 3px var(--shadow-color); | |
transform: translateZ(0); | |
transition: all 0.3s ease; | |
} | |
.message:hover { | |
transform: translateZ(5px); | |
} | |
.chat-container { | |
height: 600px !important; | |
margin-bottom: 10px; | |
} | |
.input-container { | |
height: 70px !important; | |
display: flex; | |
align-items: center; | |
gap: 10px; | |
margin-top: 5px; | |
} | |
.input-textbox { | |
height: 70px !important; | |
border-radius: 8px !important; | |
font-size: 1.1em !important; | |
padding: 10px 15px !important; | |
display: flex !important; | |
align-items: flex-start !important; /* ํ ์คํธ ์ ๋ ฅ ์์น๋ฅผ ์๋ก ์กฐ์ */ | |
} | |
.input-textbox textarea { | |
padding-top: 5px !important; /* ํ ์คํธ ์๋จ ์ฌ๋ฐฑ ์กฐ์ */ | |
} | |
.send-button { | |
height: 70px !important; | |
min-width: 70px !important; | |
font-size: 1.1em !important; | |
} | |
/* ์ค์ ํจ๋ ๊ธฐ๋ณธ ์คํ์ผ */ | |
.settings-panel { | |
padding: 20px; | |
margin-top: 20px; | |
} | |
""" | |
# GPU ๋ฉ๋ชจ๋ฆฌ ๊ด๋ฆฌ ํจ์ ์์ | |
def clear_cuda_memory(): | |
if hasattr(torch.cuda, 'empty_cache'): | |
with torch.cuda.device('cuda'): | |
torch.cuda.empty_cache() | |
# ๋ชจ๋ธ ๋ก๋ ํจ์ ์์ | |
def load_model(): | |
try: | |
model = AutoModelForCausalLM.from_pretrained( | |
MODEL_ID, | |
torch_dtype=torch.bfloat16, | |
device_map="auto", | |
) | |
return model | |
except Exception as e: | |
print(f"๋ชจ๋ธ ๋ก๋ ์ค๋ฅ: {str(e)}") | |
raise | |
def stream_chat(message: str, history: list, uploaded_file, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float): | |
global model, current_file_context | |
try: | |
if model is None: | |
model = load_model() | |
print(f'message is - {message}') | |
print(f'history is - {history}') | |
# ํ์ผ ์ ๋ก๋ ์ฒ๋ฆฌ | |
file_context = "" | |
if uploaded_file and message == "ํ์ผ์ ๋ถ์ํ๊ณ ์์ต๋๋ค...": | |
try: | |
content, file_type = read_uploaded_file(uploaded_file) | |
if content: | |
file_analysis = analyze_file_content(content, file_type) | |
file_context = f"\n\n๐ ํ์ผ ๋ถ์ ๊ฒฐ๊ณผ:\n{file_analysis}\n\nํ์ผ ๋ด์ฉ:\n```\n{content}\n```" | |
current_file_context = file_context # ํ์ผ ์ปจํ ์คํธ ์ ์ฅ | |
message = "์ ๋ก๋๋ ํ์ผ์ ๋ถ์ํด์ฃผ์ธ์." | |
except Exception as e: | |
print(f"ํ์ผ ๋ถ์ ์ค๋ฅ: {str(e)}") | |
file_context = f"\n\nโ ํ์ผ ๋ถ์ ์ค ์ค๋ฅ๊ฐ ๋ฐ์ํ์ต๋๋ค: {str(e)}" | |
elif current_file_context: # ์ ์ฅ๋ ํ์ผ ์ปจํ ์คํธ๊ฐ ์์ผ๋ฉด ์ฌ์ฉ | |
file_context = current_file_context | |
# ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋ ๋ชจ๋ํฐ๋ง | |
if torch.cuda.is_available(): | |
print(f"CUDA ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋: {torch.cuda.memory_allocated() / 1024**2:.2f} MB") | |
# ๋ํ ํ์คํ ๋ฆฌ๊ฐ ๋๋ฌด ๊ธธ๋ฉด ์๋ผ๋ด๊ธฐ | |
max_history_length = 10 # ์ต๋ ํ์คํ ๋ฆฌ ๊ธธ์ด ์ค์ | |
if len(history) > max_history_length: | |
history = history[-max_history_length:] | |
# ๊ด๋ จ ์ปจํ ์คํธ ์ฐพ๊ธฐ | |
try: | |
relevant_contexts = find_relevant_context(message) | |
wiki_context = "\n\n๊ด๋ จ ์ํคํผ๋์ ์ ๋ณด:\n" | |
for ctx in relevant_contexts: | |
wiki_context += f"Q: {ctx['question']}\nA: {ctx['answer']}\n์ ์ฌ๋: {ctx['similarity']:.3f}\n\n" | |
except Exception as e: | |
print(f"์ปจํ ์คํธ ๊ฒ์ ์ค๋ฅ: {str(e)}") | |
wiki_context = "" | |
# ๋ํ ํ์คํ ๋ฆฌ ๊ตฌ์ฑ | |
conversation = [] | |
for prompt, answer in history: | |
conversation.extend([ | |
{"role": "user", "content": prompt}, | |
{"role": "assistant", "content": answer} | |
]) | |
# ์ต์ข ํ๋กฌํํธ ๊ตฌ์ฑ | |
final_message = file_context + wiki_context + "\nํ์ฌ ์ง๋ฌธ: " + message | |
conversation.append({"role": "user", "content": final_message}) | |
# ํ ํฐ ์ ์ ํ | |
input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True) | |
max_length = 4096 # ๋๋ ๋ชจ๋ธ์ ์ต๋ ์ปจํ ์คํธ ๊ธธ์ด | |
if len(input_ids.split()) > max_length: | |
# ์ปจํ ์คํธ๊ฐ ๋๋ฌด ๊ธธ๋ฉด ์๋ผ๋ด๊ธฐ | |
input_ids = " ".join(input_ids.split()[-max_length:]) | |
inputs = tokenizer(input_ids, return_tensors="pt").to("cuda") | |
# ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋ ์ฒดํฌ | |
if torch.cuda.is_available(): | |
print(f"์ ๋ ฅ ํ ์ ์์ฑ ํ CUDA ๋ฉ๋ชจ๋ฆฌ: {torch.cuda.memory_allocated() / 1024**2:.2f} MB") | |
streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True) | |
generate_kwargs = dict( | |
inputs, | |
streamer=streamer, | |
top_k=top_k, | |
top_p=top_p, | |
repetition_penalty=penalty, | |
max_new_tokens=min(max_new_tokens, 2048), # ์ต๋ ํ ํฐ ์ ์ ํ | |
do_sample=True, | |
temperature=temperature, | |
eos_token_id=[255001], | |
) | |
# ์์ฑ ์์ ์ ๋ฉ๋ชจ๋ฆฌ ์ ๋ฆฌ | |
clear_cuda_memory() | |
thread = Thread(target=model.generate, kwargs=generate_kwargs) | |
thread.start() | |
buffer = "" | |
for new_text in streamer: | |
buffer += new_text | |
yield "", history + [[message, buffer]] | |
# ์์ฑ ์๋ฃ ํ ๋ฉ๋ชจ๋ฆฌ ์ ๋ฆฌ | |
clear_cuda_memory() | |
except Exception as e: | |
error_message = f"์ค๋ฅ๊ฐ ๋ฐ์ํ์ต๋๋ค: {str(e)}" | |
print(f"Stream chat ์ค๋ฅ: {error_message}") | |
# ์ค๋ฅ ๋ฐ์ ์์๋ ๋ฉ๋ชจ๋ฆฌ ์ ๋ฆฌ | |
clear_cuda_memory() | |
yield "", history + [[message, error_message]] | |
def create_demo(): | |
with gr.Blocks(css=CSS) as demo: | |
with gr.Column(elem_classes="markdown-style"): | |
gr.Markdown(""" | |
# ๐ค OnDevice AI RAG | |
#### ๐ RAG: ๋ฐ์ดํฐ ๋ถ์/ํ์ต ํ์ผ ์ ๋ก๋(TXT, CSV, PDF, Parquet ํ์ผ) | |
""") | |
chatbot = gr.Chatbot( | |
value=[], | |
height=600, | |
label="GiniGEN AI Assistant", | |
elem_classes="chat-container" | |
) | |
with gr.Row(elem_classes="input-container"): | |
with gr.Column(scale=1, min_width=70): | |
file_upload = gr.File( | |
type="filepath", | |
elem_classes="file-upload-icon", | |
scale=1, | |
container=True, | |
interactive=True, | |
show_label=False | |
) | |
with gr.Column(scale=3): | |
msg = gr.Textbox( | |
show_label=False, | |
placeholder="๋ฉ์์ง๋ฅผ ์ ๋ ฅํ์ธ์... ๐ญ", | |
container=False, | |
elem_classes="input-textbox", | |
scale=1 | |
) | |
with gr.Column(scale=1, min_width=70): | |
send = gr.Button( | |
"์ ์ก", | |
elem_classes="send-button custom-button", | |
scale=1 | |
) | |
with gr.Column(scale=1, min_width=70): | |
clear = gr.Button( | |
"Clear", | |
elem_classes="clear-button custom-button", | |
scale=1 | |
) | |
with gr.Accordion("๐ฎ ๊ณ ๊ธ ์ค์ ", open=False): | |
with gr.Row(): | |
with gr.Column(scale=1): | |
temperature = gr.Slider( | |
minimum=0, maximum=1, step=0.1, value=0.8, | |
label="์ฐฝ์์ฑ ์์ค ๐จ" | |
) | |
max_new_tokens = gr.Slider( | |
minimum=128, maximum=8000, step=1, value=4000, | |
label="์ต๋ ํ ํฐ ์ ๐" | |
) | |
with gr.Column(scale=1): | |
top_p = gr.Slider( | |
minimum=0.0, maximum=1.0, step=0.1, value=0.8, | |
label="๋ค์์ฑ ์กฐ์ ๐ฏ" | |
) | |
top_k = gr.Slider( | |
minimum=1, maximum=20, step=1, value=20, | |
label="์ ํ ๋ฒ์ ๐" | |
) | |
penalty = gr.Slider( | |
minimum=0.0, maximum=2.0, step=0.1, value=1.0, | |
label="๋ฐ๋ณต ์ต์ ๐" | |
) | |
gr.Examples( | |
examples=[ | |
["๋ค์ ์ฝ๋์ ๋ฌธ์ ์ ์ ์ฐพ์๋ด๊ณ ๊ฐ์ ๋ ๋ฒ์ ์ ์ ์ํด์ฃผ์ธ์:\ndef fibonacci(n):\n if n <= 1: return n\n return fibonacci(n-1) + fibonacci(n-2)"], | |
["์ฃผ์ด์ง ๋ฐ์ดํฐ๋ฅผ ๋ถ์ํ๊ณ ์ธ์ฌ์ดํธ๋ฅผ ๋์ถํด์ฃผ์ธ์:\n์ฐ๋๋ณ ๋งค์ถ์ก(์ต์)\n2019: 1200\n2020: 980\n2021: 1450\n2022: 2100\n2023: 1890"], | |
["๋ค์ ์ํ ๋ฌธ์ ๋ฅผ ๋จ๊ณ๋ณ๋ก ์์ธํ ํ์ดํด์ฃผ์ธ์: 'ํ ์์ ๋์ด๊ฐ ๊ทธ ์์ ๋ด์ ํ๋ ์ ์ฌ๊ฐํ ๋์ด์ 2๋ฐฐ์ผ ๋, ์์ ๋ฐ์ง๋ฆ๊ณผ ์ ์ฌ๊ฐํ์ ํ ๋ณ์ ๊ธธ์ด์ ๊ด๊ณ๋ฅผ ๊ตฌํ์์ค.'"], | |
["๋ค์ ๋ง์ผํ ์บ ํ์ธ์ ROI๋ฅผ ๋ถ์ํ๊ณ ๊ฐ์ ๋ฐฉ์์ ์ ์ํด์ฃผ์ธ์:\n์ด ๋น์ฉ: 5000๋ง์\n๋๋ฌ์ ์: 100๋ง๋ช \nํด๋ฆญ๋ฅ : 2.3%\n์ ํ์จ: 0.8%\nํ๊ท ๊ตฌ๋งค์ก: 35,000์"], | |
], | |
inputs=msg | |
) | |
# Clear ํจ์ ์ ์ | |
def clear_conversation(): | |
global current_file_context | |
current_file_context = None | |
return [], None, "์๋ก์ด ๋ํ๋ฅผ ์์ํ์ธ์..." | |
# ์ด๋ฒคํธ ๋ฐ์ธ๋ฉ | |
msg.submit( | |
stream_chat, | |
inputs=[msg, chatbot, file_upload, temperature, max_new_tokens, top_p, top_k, penalty], | |
outputs=[msg, chatbot] | |
) | |
send.click( | |
stream_chat, | |
inputs=[msg, chatbot, file_upload, temperature, max_new_tokens, top_p, top_k, penalty], | |
outputs=[msg, chatbot] | |
) | |
file_upload.change( | |
fn=init_msg, | |
outputs=msg, | |
queue=False | |
).then( | |
fn=stream_chat, | |
inputs=[msg, chatbot, file_upload, temperature, max_new_tokens, top_p, top_k, penalty], | |
outputs=[msg, chatbot], | |
queue=True | |
) | |
# Clear ๋ฒํผ ์ด๋ฒคํธ ๋ฐ์ธ๋ฉ | |
clear.click( | |
fn=clear_conversation, | |
outputs=[chatbot, file_upload, msg], | |
queue=False | |
) | |
return demo | |
if __name__ == "__main__": | |
demo = create_demo() | |
demo.launch() |