File size: 1,524 Bytes
3760948
 
 
 
6ea9abd
3760948
6ea9abd
 
 
 
3760948
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6ea9abd
3760948
 
 
 
 
 
 
 
 
ec9a090
3760948
 
 
 
 
 
 
6ea9abd
03fd821
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
43
44
45
46
47
48
49
50
51
52
53
54
55
from fastapi import FastAPI
import os
import json
import google.generativeai as genai
from pydantic import BaseModel, validator


class Item(BaseModel):
    text: str = "sddddddddddd"
    
app = FastAPI()

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['text']}",
        "Category: ",
    ]

    response = model.generate_content(prompt_parts)

    json_response = json.loads(
        response.text[response.text.find("{") : response.text.rfind("}") + 1]
    )

    return json_response['category']

    
@app.get("/")
async def root():
 return {"Text Emotion Classification":"Version 1.5 'Text'"}

@app.post("/classify/")
def read_user(js: Item):
    return classify_msg(js.dict())