Spaces:

GueuleDange
/

Test

Sleeping

App Files Files Community

GueuleDange commited on Mar 28

Commit

3796eb5

verified ·

1 Parent(s): e41d8d9

Update app.py

Browse files

Files changed (1) hide show

app.py +35 -60

app.py CHANGED Viewed

@@ -1,26 +1,28 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse, StreamingResponse
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import asyncio
 app = FastAPI()
-# Chargement du modèle
 model_name = "microsoft/Phi-3.5-mini-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float16,
-    device_map="auto"
 )
-# Fonction de génération avec streaming
-async def generate_stream(prompt: str):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    for _ in range(512):  # Limite de tokens
         outputs = model.generate(
             **inputs,
             max_new_tokens=1,
@@ -30,61 +32,34 @@ async def generate_stream(prompt: str):
         )
         new_token = tokenizer.decode(outputs[0][-1], skip_special_tokens=True)
         yield f"data: {new_token}\n\n"
-        await asyncio.sleep(0.05)
         inputs = {"input_ids": outputs}
-# Interface Gradio standard
-def generate_text(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=512)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-# Page web de streaming
 @app.get("/", response_class=HTMLResponse)
-async def web_interface(request: Request):
-    return """
-    <!DOCTYPE html>
-    <html>
-    <head>
-        <title>Chat Streaming</title>
-        <script>
-            async function startStream() {
-                const prompt = document.getElementById("prompt").value;
-                const output = document.getElementById("output");
-                output.innerHTML = "";
-                const eventSource = new EventSource(`/stream?prompt=${encodeURIComponent(prompt)}`);
-                eventSource.onmessage = (event) => {
-                    output.innerHTML += event.data;
-                    output.scrollTop = output.scrollHeight;
-                };
-            }
-        </script>
-    </head>
-    <body>
-        <h1>Chat en temps réel</h1>
-        <textarea id="prompt" rows="4"></textarea>
-        <button onclick="startStream()">Envoyer</button>
-        <div id="output" style="white-space: pre-wrap; margin-top: 20px;"></div>
-    </body>
-    </html>
-    """
-# Endpoint de streaming
 @app.get("/stream")
 async def stream_response(prompt: str):
-    return StreamingResponse(
-        generate_stream(prompt),
-        media_type="text/event-stream"
-    )
-# Interface Gradio (accessible via /gradio)
-demo = gr.Interface(
-    fn=generate_text,
-    inputs="text",
-    outputs="text",
-    title="Phi-3 Chat"
-)
-app = gr.mount_gradio_app(app, demo, path="/gradio")

+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi.templating import Jinja2Templates
 import torch
 import asyncio
+from transformers import AutoTokenizer, AutoModelForCausalLM
 app = FastAPI()
+templates = Jinja2Templates(directory="templates")
+# Configuration du modèle (optimisé pour 2000 tokens)
 model_name = "microsoft/Phi-3.5-mini-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float16,
+    device_map="auto",
+    low_cpu_mem_usage=True  # Critique pour les longs contextes
 )
+async def generate_stream(prompt: str, max_tokens: int = 2000):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    generated_count = 0
+    while generated_count < max_tokens:
         outputs = model.generate(
             **inputs,
             max_new_tokens=1,
         )
         new_token = tokenizer.decode(outputs[0][-1], skip_special_tokens=True)
         yield f"data: {new_token}\n\n"
+        generated_count += 1
+        # Optimisation mémoire
+        if generated_count % 50 == 0:
+            await asyncio.sleep(0.01)  # Réduit la pression sur le GPU
+            torch.cuda.empty_cache()  # Nettoyage mémoire
         inputs = {"input_ids": outputs}
 @app.get("/", response_class=HTMLResponse)
+async def chat_page(request: Request):
+    return templates.TemplateResponse("stream.html", {"request": request})
 @app.get("/stream")
 async def stream_response(prompt: str):
+    try:
+        return StreamingResponse(
+            generate_stream(prompt),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no"  # Critique pour les streams longs
+            }
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)