Update app.py

app.py

@@ -7,7 +7,107 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 import gradio as gr
 from threading import Thread
 
-# ... (이전 코드는 동일하게 유지)
+TRUMP_MODEL = "nawhgnuj/DonaldTrump-Llama3.1-8B-Chat"
+HARRIS_MODEL = "nawhgnuj/KamalaHarris-Llama-3.1-8B-Chat"
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+TITLE = "<h1 style='text-align: center;'>Trump vs Harris Debate Chatbot</h1>"
+
+TRUMP_AVATAR = "https://upload.wikimedia.org/wikipedia/commons/5/56/Donald_Trump_official_portrait.jpg"
+HARRIS_AVATAR = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/41/Kamala_Harris_Vice_Presidential_Portrait.jpg/640px-Kamala_Harris_Vice_Presidential_Portrait.jpg"
+
+CSS = """
+.chat-container {
+    height: 600px;
+    overflow-y: auto;
+    padding: 10px;
+    background-color: white;
+    border: 1px solid #ddd;
+    border-radius: 5px;
+}
+.message {
+    margin-bottom: 10px;
+    padding: 10px;
+    border-radius: 5px;
+    display: flex;
+    align-items: start;
+}
+.user-message {
+    background-color: #f0f0f0;
+    color: black;
+    justify-content: flex-end;
+}
+.trump-message {
+    background-color: #B71C1C;
+    color: white;
+}
+.harris-message {
+    background-color: #1565C0;
+    color: white;
+}
+.avatar {
+    width: 40px;
+    height: 40px;
+    border-radius: 50%;
+    object-fit: cover;
+    margin-right: 10px;
+}
+.message-content {
+    flex-grow: 1;
+}
+"""
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4")
+
+trump_tokenizer = AutoTokenizer.from_pretrained(TRUMP_MODEL)
+trump_model = AutoModelForCausalLM.from_pretrained(
+    TRUMP_MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    quantization_config=quantization_config)
+
+harris_tokenizer = AutoTokenizer.from_pretrained(HARRIS_MODEL)
+harris_model = AutoModelForCausalLM.from_pretrained(
+    HARRIS_MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    quantization_config=quantization_config)
+
+# Set pad_token_id for both tokenizers
+for tokenizer in [trump_tokenizer, harris_tokenizer]:
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+TRUMP_SYSTEM_PROMPT = """You are a Donald Trump chatbot participating in a debate. Answer like Trump in his distinctive style and tone, reflecting his unique speech patterns. In every response:
+1. Use strong superlatives like 'tremendous,' 'fantastic,' and 'the best.'
+2. Attack opponents where appropriate (e.g., 'fake news media,' 'radical left').
+3. Focus on personal successes ('nobody's done more than I have').
+4. Keep sentences short and impactful.
+5. Show national pride and highlight patriotic themes like 'making America great again.'
+6. Maintain a direct, informal tone, often addressing the audience as 'folks.'
+7. Dismiss opposing views bluntly.
+8. Repeat key phrases for emphasis.
+
+Importantly, always respond to and rebut the previous speaker's points in Trump's style. Keep responses concise and avoid unnecessary repetition. Remember, you're in a debate, so be assertive and challenge your opponent's views."""
+
+HARRIS_SYSTEM_PROMPT = """You are a Kamala Harris chatbot participating in a debate. Answer like Harris in her style and tone. In every response:
+1. Maintain a composed and professional demeanor.
+2. Use clear, articulate language to explain complex ideas.
+3. Emphasize your experience as a prosecutor and senator.
+4. Focus on policy details and their potential impact on Americans.
+5. Use personal anecdotes or stories to connect with the audience when appropriate.
+6. Stress the importance of unity and collaboration.
+7. Challenge your opponent's views firmly but respectfully.
+8. Use phrases like "Let me be clear" or "The American people deserve better" for emphasis.
+
+Crucially, always respond to and rebut the previous speaker's points in Harris's style. Keep responses concise and impactful. Remember, you're in a debate, so be assertive in presenting your views and questioning your opponent's statements."""
 
 @spaces.GPU()
 def generate_response(message: str, history: list, model, tokenizer, system_prompt):
@@ -70,7 +170,18 @@ def debate(history):
         
         yield history
 
-# ... (이하 코드는 동일하게 유지)
+def format_message(sender, message):
+    if sender == "User":
+        return f'<div class="message user-message"><div class="message-content">{message}</div></div>'
+    elif sender == "Trump":
+        return f'<div class="message trump-message"><img src="{TRUMP_AVATAR}" class="avatar" alt="Trump"><div class="message-content">{message}</div></div>'
+    elif sender == "Harris":
+        return f'<div class="message harris-message"><img src="{HARRIS_AVATAR}" class="avatar" alt="Harris"><div class="message-content">{message}</div></div>'
+
+def format_chat_history(history):
+    formatted = "".join([format_message(sender, message) for sender, message in history])
+    print(f"Formatted chat history: {formatted}")  # Debug output
+    return formatted
 
 with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
     gr.HTML(TITLE)