Update app.py

app.py

@@ -1,18 +1,20 @@
+import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM 
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_name = "beomi/gemma-ko-2b"  # 혹은 더 가벼운 모델 선택
+model_name = "beomi/gemma-ko-2b"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)  
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-prompt = "인공지능이란 무엇인가?"
-with torch.no_grad():
-    tokens = tokenizer.encode(prompt, return_tensors='pt').to(device)
-    gen_tokens = model.generate(tokens, do_sample=True, temperature=0.8, max_length=64)
-    generated = tokenizer.decode(gen_tokens[0], skip_special_tokens=True)
+def chatbot(prompt):
+    with torch.no_grad():
+        tokens = tokenizer(prompt, return_tensors='pt').to(device)
+        gen_tokens = model.generate(tokens, do_sample=True, temperature=0.8, max_length=64)
+        return tokenizer.decode(gen_tokens[0], skip_special_tokens=True)
 
-print(generated)
+iface = gr.Interface(fn=chatbot, inputs="text", outputs="text")
+iface.launch()