a

Dockerfile

@@ -1,16 +1,11 @@
-FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04
+FROM python:3.10
 
-# Installer les dépendances système
-RUN apt-get update && apt-get install -y \
-    python3 python3-pip git && \
-    ln -s /usr/bin/python3 /usr/bin/python
+WORKDIR /app
 
-# Installer VLLM
-RUN pip install --upgrade pip
-RUN pip install vllm
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
 
-# Ajouter ton script
-COPY app.py /app.py
+COPY . .
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
 
-# Lancer le serveur VLLM
-CMD ["python", "/app.py"]

app.py

@@ -1,12 +1,24 @@
-import subprocess
-
-# Exemple avec un modèle léger (à adapter)
-model = "TheBloke/Mistral-7B-Instruct-v0.1-GGUF"
-
-# Lancement de VLLM en mode API
-subprocess.run([
-    "python3", "-m", "vllm.entrypoints.api_server",
-    "--model", model,
-    "--host", "0.0.0.0",
-    "--port", "7860"
-])
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+app = FastAPI()
+
+# Charger le modèle
+model_name = "google/medgemma-4b-pt"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+
+# Requête attendue
+class Request(BaseModel):
+    prompt: str
+
+@app.post("/generate")
+def generate(request: Request):
+    inputs = tokenizer(request.prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=100)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": result}
+

requirement.txt

@@ -1 +1,5 @@
-vllm
+fastapi
+uvicorn
+transformers
+torch
+