Update app.py

app.py

@@ -1,39 +1,44 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, GemmaTokenizer
 import torch
 import os
+import gradio as gr
 
-# 모델 및 토크나이저 로드
-model_id = "kimhyunwoo/gemma2-ko-dialogue-lora-fp16"
-model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
-# AutoTokenizer 대신 직접 GemmaTokenizer를 로드합니다.
-tokenizer = GemmaTokenizer.from_pretrained(model_id, trust_remote_code=True)
+try:
+    # 모델 및 토크나이저 로드
+    model_id = "kimhyunwoo/gemma2-ko-dialogue-lora-fp16"
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
+    # AutoTokenizer 대신 직접 GemmaTokenizer를 로드합니다.
+    tokenizer = GemmaTokenizer.from_pretrained(model_id, trust_remote_code=True)
 
-# CPU 환경에서 모델을 로드하고 사용하는 경우 torch.float32를 명시적으로 사용할 수 있습니다.
-model = model.to(torch.float32)
-model.eval()
+    # CPU 환경에서 모델을 로드하고 사용하는 경우 torch.float32를 명시적으로 사용할 수 있습니다.
+    model = model.to(torch.float32)
+    model.eval()
 
-# 텍스트 생성 함수
-def generate_text(text, max_length=50, do_sample=True, temperature=1.0):
-    inputs = tokenizer(text, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=max_length, do_sample=do_sample, temperature=temperature)
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return generated_text
+    # 텍스트 생성 함수
+    def generate_text(text, max_length=50, do_sample=True, temperature=1.0):
+        inputs = tokenizer(text, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_length=max_length, do_sample=do_sample, temperature=temperature)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return generated_text
 
-# Gradio 인터페이스 생성
-import gradio as gr
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox(lines=5, placeholder="텍스트를 입력하세요..."),
-        gr.Slider(minimum=10, maximum=200, step=1, value=50, label="최대 길이"),
-        gr.Checkbox(label="샘플링"),
-        gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="온도"),
-    ],
-    outputs=gr.Textbox(lines=5, label="생성된 텍스트"),
-    title="Gemma 2 Text Generator",
-    description="Fine-tuned된 Gemma 2 모델을 사용하여 텍스트를 생성합니다.",
-)
+    # Gradio 인터페이스 생성
+    iface = gr.Interface(
+        fn=generate_text,
+        inputs=[
+            gr.Textbox(lines=5, placeholder="텍스트를 입력하세요..."),
+            gr.Slider(minimum=10, maximum=200, step=1, value=50, label="최대 길이"),
+            gr.Checkbox(label="샘플링"),
+            gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="온도"),
+        ],
+        outputs=gr.Textbox(lines=5, label="생성된 텍스트"),
+        title="Gemma 2 Text Generator",
+        description="Fine-tuned된 Gemma 2 모델을 사용하여 텍스트를 생성합니다.",
+    )
 
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    if __name__ == "__main__":
+        iface.launch(server_name="0.0.0.0", server_port=7860)
+except ImportError as e:
+    print(f"ImportError 발생: {e}")
+except Exception as e:
+    print(f"예상치 못한 오류 발생: {e}")