EU_AI_ACT / app.py
Komal-patra's picture
updated
29f9dea verified
raw
history blame
3.9 kB
import os
import gradio as gr
from huggingface_hub import login
from transformers import AutoModelForSeq2SeqLM, T5Tokenizer
from peft import PeftModel, PeftConfig
# Hugging Face login
token = os.environ.get("token")
if token:
login(token)
print("Login is successful")
else:
print("Token not found. Please set your token in the environment variables.")
# Model and tokenizer setup
MODEL_NAME = "google/flan-t5-base"
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, use_auth_token=token)
config = PeftConfig.from_pretrained("Komal-patra/results")
base_model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
model = PeftModel.from_pretrained(base_model, "Komal-patra/results")
# Text generation function
def generate_text(prompt, max_length=150):
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(
input_ids=inputs["input_ids"],
max_length=max_length,
num_beams=1,
repetition_penalty=2.2
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_text
# Custom CSS for the UI
custom_css = """
.message.pending {
background: #A8C4D6;
}
/* Response message */
.message.bot.svelte-1s78gfg.message-bubble-border {
border-color: #266B99;
}
/* User message */
.message.user.svelte-1s78gfg.message-bubble-border {
background: #9DDDF9;
border-color: #9DDDF9;
}
/* For both user and response message as per the document */
span.md.svelte-8tpqd2.chatbot.prose p {
color: #266B99;
}
/* Chatbot container */
.gradio-container {
background: #1c1c1c; /* Dark background */
color: white; /* Light text color */
}
/* RED (Hex: #DB1616) for action buttons and links only */
.clear-btn {
background: #DB1616;
color: white;
}
/* Primary colors are set to be used for all sorts */
.submit-btn {
background: #266B99;
color: white;
}
/* Add icons to messages */
.message.user.svelte-1s78gfg {
display: flex;
align-items: center;
}
.message.user.svelte-1s78gfg:before {
content: url('file=Komal-patra/EU_AI_ACT/user icon.jpeg');
margin-right: 8px;
}
.message.bot.svelte-1s78gfg {
display: flex;
align-items: center;
}
.message.bot.svelte-1s78gfg:before {
content: url('file=Komal-patra/EU_AI_ACT/orcawise image.png');
margin-right: 8px;
}
/* Enable scrolling for the chatbot messages */
.chatbot .messages {
max-height: 500px; /* Adjust as needed */
overflow-y: auto;
}
"""
# Gradio interface setup
with gr.Blocks(css=custom_css) as demo:
gr.Markdown("<h1>Ask a question about the EU AI Act</h1>")
chatbot = gr.Chatbot()
msg = gr.Textbox(placeholder="Ask your question...", show_label=False) # Add placeholder text
submit_button = gr.Button("Submit", elem_classes="submit-btn")
clear = gr.Button("Clear", elem_classes="clear-btn")
# Function to handle user input
def user(user_message, history):
return "", history + [[user_message, None]]
# Function to handle bot response
def bot(history):
if len(history) == 1: # Check if it's the first interaction
bot_message = "Hi there! How can I help you today?"
history[-1][1] = bot_message # Add welcome message to history
else:
history[-1][1] = "" # Clear the last bot message
previous_message = history[-1][0] # Access the previous user message
bot_message = generate_text(previous_message) # Generate response based on previous message
history[-1][1] = bot_message # Update the last bot message
return history
submit_button.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
bot, chatbot, chatbot
)
msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
bot, chatbot, chatbot
)
clear.click(lambda: None, None, chatbot, queue=False)
demo.launch()