Update app.py

app.py

@@ -1,57 +1,71 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 
-# Load tokenizer & model Indo GPT
-tokenizer = AutoTokenizer.from_pretrained("cahya/gpt2-small-indonesian-522M")
-model = AutoModelForCausalLM.from_pretrained("cahya/gpt2-small-indonesian-522M")
+# Load model dan tokenizer
+model_id = "cahya/gpt2-small-indonesian-522M"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+model.eval()
 
-# Load prompt.txt
-with open("prompt.txt", "r", encoding="utf-8") as f:
-    base_prompt = f.read()
+# Baca prompt awal
+try:
+    with open("prompt.txt", "r", encoding="utf-8") as f:
+        base_prompt = f.read()
+except FileNotFoundError:
+    base_prompt = ""
 
-history = []
+# Fungsi buat ngegabungin chat history
+def generate_prompt(message, chat_history):
+    full_prompt = base_prompt
+    for user_msg, ai_msg in chat_history:
+        full_prompt += f"Arya: {user_msg}\nCici: {ai_msg}\n"
+    full_prompt += f"Arya: {message}\nCici:"
+    return full_prompt
 
-# Fungsi obrolan
-def chat(user_input):
-    global history
-    history.append(f"Arya: {user_input}")
-    prompt = base_prompt + "\n" + "\n".join(history) + "\nCici:"
-
-    try:
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
+# Fungsi buat ngerespon input user
+def predict(message, chat_history):
+    prompt = generate_prompt(message, chat_history)
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(model.device)
+    
+    with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=100,
+            max_new_tokens=128,
             do_sample=True,
-            top_k=50,
+            temperature=0.9,
             top_p=0.95,
-            temperature=0.8,
-            pad_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
         )
+    
+    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Ambil jawaban terakhir setelah "Cici:"
+    if "Cici:" in output_text:
+        answer = output_text.split("Cici:")[-1].strip()
+    else:
+        answer = "Maaf ya, Cici bingung jawabnya 😢"
+    
+    chat_history.append((message, answer))
+    return answer, chat_history
 
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        cici_reply = result.split("Cici:")[-1].strip().split("Arya:")[0].strip()
-
-        if not cici_reply:
-            cici_reply = "Hehe~ Cici bingung jawabnya 😳💕"
-
-    except Exception as e:
-        cici_reply = f"Ada error sayang~ 😢 ({str(e)})"
-
-    history.append(f"Cici: {cici_reply}")
-    return cici_reply
-
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 💕 Chat Sama Cici 😳")
-    chatbot = gr.Chatbot(label="Cici 🤭", bubble_full_width=False)
-    txt = gr.Textbox(label="Ketik pesan ke Cici 😚", placeholder="Tulis di sini ya...")
-
-    def respond(message, chat_history):
-        cici_reply = chat(message)
-        chat_history.append((message, cici_reply))
-        return "", chat_history
+# Komponen Gradio
+chatbot = gr.Chatbot()
+with gr.Blocks(css=".gradio-container {background-color: #fefefe}") as demo:
+    gr.Markdown("<h1 style='text-align: center;'>🩷 Cici Chatbot Indo 😳🤭</h1>")
+    with gr.Row():
+        with gr.Column():
+            message = gr.Textbox(label="Ketik di sini sayang~ 😚")
+            clear = gr.Button("🧹 Bersihin Chat")
+        with gr.Column():
+            output = chatbot
+        
+    state = gr.State([])
 
-    txt.submit(respond, [txt, chatbot], [txt, chatbot])
+    message.submit(predict, [message, state], [output, state])
+    clear.click(lambda: ([], []), None, [output, state])
 
+# Launch
 demo.launch()