change model to llm write

app.py

@@ -64,27 +64,29 @@ Merci !
 
 """
 class PredictionRequest(BaseModel):
-    history: List[Tuple[str, str]] = []
-    prompt: str = ""
+    history: list
+    prompt: str
     max_length: int = 128000
     top_p: float = 0.8
     temperature: float = 0.6
-@app.post("/generate/")
-async def predict(request: PredictionRequest):
-    history = default_prompt
-    prompt = request.prompt
-    max_length = request.max_length
-    top_p = request.top_p
-    temperature = request.temperature
 
+class PredictionResponse(BaseModel):
+    history: list
+
+@app.get("/")
+async def home():
+    return 'STN BIG DATA'
+@app.post("/predict", response_model=PredictionResponse)
+async def predict(request: PredictionRequest):
     stop = StopOnTokens()
     messages = []
-    if prompt:
-        messages.append({"role": "system", "content": prompt})
-    for idx, (user_msg, model_msg) in enumerate(history):
-        if prompt and idx == 0:
+    query = ""
+
+    # Préparer les messages sans inclure le prompt par défaut dans l'historique renvoyé
+    for idx, (user_msg, model_msg) in enumerate(request.history):
+        if idx == 0 and request.prompt:  # Ignorer le prompt par défaut dans l'historique
             continue
-        if idx == len(history) - 1 and not model_msg:
+        if idx == len(request.history) - 1 and not model_msg:
             query = user_msg
             break
         if user_msg:
@@ -92,35 +94,35 @@ async def predict(request: PredictionRequest):
         if model_msg:
             messages.append({"role": "assistant", "content": model_msg})
 
+    # Inclure le prompt uniquement pour la génération, mais pas dans l'historique
     model_inputs = tokenizer.build_chat_input(query, history=messages, role='user').input_ids.to(
-        next(model.parameters()).device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=600, skip_prompt=True)
+        next(model.parameters()).device
+    )
+
     eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
                     tokenizer.get_command("<|observation|>")]
+
     generate_kwargs = {
         "input_ids": model_inputs,
-        "streamer": streamer,
-        "max_new_tokens": max_length,
+        "max_new_tokens": request.max_length,
         "do_sample": True,
-        "top_p": top_p,
-        "temperature": temperature,
+        "top_p": request.top_p,
+        "temperature": request.temperature,
         "stopping_criteria": StoppingCriteriaList([stop]),
         "repetition_penalty": 1,
         "eos_token_id": eos_token_id,
     }
 
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+    # Générer le texte
+    output = model.generate(**generate_kwargs)
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Ajouter le texte généré à l'historique
+    request.history[-1][1] += generated_text
 
-    generated_text = ""
-    for new_token in streamer:
-        if new_token and '<|user|>' in new_token:
-            new_token = new_token.split('<|user|>')[0]
-        if new_token:
-            generated_text += new_token
-            history[-1][1] = generated_text
+    # Retourner l'historique sans le prompt
+    return PredictionResponse(history=request.history)
 
-    return {"history": history}
 if __name__ == "__main__":
     uvicorn.run("app:app",reload=True)