Spaces:
Running
Running
File size: 1,300 Bytes
856414b 2a5ffea 856414b 7b28bf8 856414b 90bd4f9 856414b 8d40135 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import os
import json
import gradio as gr
import google.generativeai as genai
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
genai.configure(api_key=GOOGLE_API_KEY)
# Set up the model
generation_config = {
"temperature": 0.9,
"top_p": 1,
"top_k": 1,
"max_output_tokens": 2048,
}
model = genai.GenerativeModel(
model_name="gemini-pro",
generation_config=generation_config,
)
task_description = " You need to classify each message you receive among the following categories: 'admiration','amusement','anger','annoyance','approval','caring','confusion','curiosity','desire','disappointment','disapproval','disgust','embarrassment','excitement','fear','gratitude','grief','joy','love','nervousness', 'optimism', 'pride', 'realization', 'relief', 'remorse', 'sadness', 'surprise', 'neutral'<div>The output must be in JSON format</div>"
def classify_msg(message):
prompt_parts = [
task_description,
f"Message: {message}",
"Category: ",
]
response = model.generate_content(prompt_parts)
json_response = json.loads(
response.text[response.text.find("{") : response.text.rfind("}") + 1]
)
return gr.Label(json_response['category'])
iface = gr.Interface(fn=classify_msg, inputs="text", outputs="text")
iface.launch() |