version 2

app.py

@@ -1,101 +1,63 @@
-#!/usr/bin/env python
-
-import os
-from collections.abc import Iterator
-from threading import Thread
-
-import gradio as gr
-import spaces
+from fastapi import FastAPI, Request
+from fastapi.responses import HTMLResponse, JSONResponse
+import uvicorn
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 
-DESCRIPTION = """
-<h1 style="color:black;">Mistral-7B v0.3</h1>
-"""
-
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
-
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+app = FastAPI()
 
+# Chargement du modèle uniquement si CUDA est disponible
 if torch.cuda.is_available():
     model_id = "mistralai/Mistral-7B-Instruct-v0.3"
     model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
     tokenizer = AutoTokenizer.from_pretrained(model_id)
+else:
+    model = None
+    tokenizer = None
 
+MAX_INPUT_TOKEN_LENGTH = 4096
 
-@spaces.GPU
-def generate(
-    message: str,
-    chat_history: list[dict],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
-) -> Iterator[str]:
-    conversation = [*chat_history, {"role": "user", "content": message}]
+def generate_response(message: str, history: list) -> str:
+    conversation = history + [{"role": "user", "content": message}]
     input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     input_ids = input_ids.to(model.device)
-
+    
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-        repetition_penalty=repetition_penalty,
-    )
+    generate_kwargs = {
+        "input_ids": input_ids,
+        "streamer": streamer,
+        "max_new_tokens": 1024,
+        "do_sample": True,
+        "top_p": 0.9,
+        "top_k": 50,
+        "temperature": 0.6,
+        "num_beams": 1,
+        "repetition_penalty": 1.2,
+    }
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
-
-    outputs = []
+    
+    response_text = ""
     for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
-
-
-# CSS pour appliquer le dégradé pastel à TOUTE la page
-custom_css = """
-html, body {
-    height: 100%;
-    margin: 0;
-    padding: 0;
-    background: linear-gradient(135deg, #FDE2E2, #E2ECFD) !important;
-}
-"""
-
-# Questions prédéfinies
-predefined_examples = [
-    ["1 - C’est quoi le consentement ? Comment savoir si ma copine a envie de moi ?"],  # noqa: RUF001
-    ["2 - C’est quoi une agression sexuelle ?"],
-    ["3 - C’est quoi un viol ?"],
-    ["4 - C’est quoi un attouchement ?"],
-    ["5 - C’est quoi un harcèlement sexuel ?"],
-    ["6 - Est ce illégal de visionner du porno ?"],
-    ["7 - Mon copain me demande un nude, dois-je le faire ?"],
-    ["8 - Mon ancien copain me menace de poster des photos de moi nue sur internet, que faire ?"],
-    [
-        "9 - Que puis-je faire si un membre de ma famille me touche d’une manière bizarre, mais que j’ai peur de parler ou de ne pas être cru ?"
-    ],
-]
-
-demo = gr.ChatInterface(
-    fn=generate,
-    type="messages",
-    description=DESCRIPTION,
-    css=custom_css,  # On applique le CSS pastel global
-    examples=predefined_examples,
-)
+        response_text += text
+    return response_text
+
+@app.post("/chat")
+async def chat_endpoint(request: Request):
+    data = await request.json()
+    message = data.get("message", "")
+    # Utilisation d'un historique vide pour simplifier
+    response_text = generate_response(message, history=[])
+    return JSONResponse({"response": response_text})
+
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    with open("index.html", "r", encoding="utf-8") as f:
+        html_content = f.read()
+    return HTMLResponse(content=html_content, status_code=200)
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    uvicorn.run(app, host="0.0.0.0", port=8000)