File size: 4,597 Bytes
291d462
 
 
 
 
 
 
 
 
 
 
 
 
 
4342296
291d462
 
 
 
 
 
 
 
 
1770fad
291d462
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import InferenceClient
import uvicorn
from paho.mqtt import client as mqtt_client
import uuid
import time



app = FastAPI()
origins = [
    "*",
    "http://localhost:3000",
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

inferenceClient = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
historyMap = {}

broker_host = "broker.emqx.io"
broker_port = 1883
broker_topic_prefix = "/diopsoft/mixtral-fastapi"
# mqttClient = None
mqttClient = mqtt_client.Client(mqtt_client.CallbackAPIVersion.VERSION1, str(uuid.uuid1()))
def connectMqttBroker():
    try:
        res = mqttClient.connect(broker_host, broker_port, 60)
        print('MQTT connection status : ' + str(res))
        mqttClient.loop_start()
    except Exception as Argument:
        # logging.getLogger().exception("Failed to connect to MQTT Broker")
        # logging.getLogger().error(e)
        traceback.print_exc()
def on_connect(client, userdata, flags, rc):
        if rc == 0:
            print("Connected to MQTT Broker!")
        else:
            raise Exception("Failed to connect to MQTT Broker, return code " + rc)
mqttClient.on_connect = on_connect
# connectMqttBroker()


class Item(BaseModel):
    is_mock: bool = False
    token: str = None
    mqttTopic: str = None
    prompt: str
    history: list = None
    # system_prompt: str = "Soyez un assistant utile, en français de préférence"
    system_prompt: str = "Soyez un assistant utile"
    temperature: float = 0.7
    max_new_tokens: int = 1048
    top_p: float = 0.9
    repetition_penalty: float = 1.1

def format_prompt(message, history):
    prompt = "<s>"
    if history:
        for hist in history:
            prompt += f"[INST] {hist['user_prompt']} [/INST]"
            prompt += f" {hist['bot_response']}</s> "
    prompt += f"[INST] {message} [/INST]"
    return prompt

def generate(item: Item):
    temperature = float(item.temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(item.top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=item.max_new_tokens,
        top_p=top_p,
        repetition_penalty=item.repetition_penalty,
        do_sample=True,
        seed=42,
    )

    if item.mqttTopic and not mqttClient.is_connected():
        connectMqttBroker()

    token = item.token
    history = None
    if token:
        history = historyMap.get(token)

    formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", history)
    if not item.is_mock:
        stream = inferenceClient.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    if not item.is_mock:
        for response in stream:
            print(response.token.text)
            output += response.token.text
            if item.mqttTopic and mqttClient.is_connected():
                # result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, response.token.text)
                result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, output)
                status = result[0]
                if status == 0:
                    print("Message sent to topic")
                else:
                    print("Failed to send message to topic")
    else:
        for i in range(15):
            text = " generated text " + str(i)
            time.sleep(0.5)
            print(text)
            output += text
            if item.mqttTopic and mqttClient.is_connected():
                # result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, text)
                result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, output)
                status = result[0]
                if status == 0:
                    print("Message sent to topic")
                else:
                    print("Failed to send message to topic")

    if item.mqttTopic and mqttClient.is_connected():
        mqttClient.unsubscribe(broker_topic_prefix+'/'+item.mqttTopic)

    if token:
        if not historyMap.get(token):
            historyMap[token] = []
        historyMap.get(token).append({ "user_prompt":item.prompt, "bot_response":output })

    return output

@app.post("/generate/")
async def generate_text(item: Item):
    return {"response": generate(item)}

@app.get("/test-api/")
async def testAPI():
    return "API V2 is running"