Robostar commited on
Commit
961ab91
·
verified ·
1 Parent(s): 4c525ad

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer
4
+ import torch
5
+
6
+ app = FastAPI()
7
+
8
+ # Charger le modèle et le tokenizer
9
+ #model_name = "mistralai/Mistral-7B-Instruct-v0.1" # Modèle Mistral 7B
10
+ #model_name = "HuggingFaceH4/zephyr-3b"
11
+ #model_name = "serkanarslan/mistral-7b-mini-ft"
12
+ model_name = "microsoft/phi-2"
13
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
14
+ model = AutoModelForCausalLM.from_pretrained(
15
+ model_name,
16
+ torch_dtype=torch.float16,
17
+ device_map="auto" # Utilise le GPU si dispo
18
+ )
19
+
20
+ # Définir le format des requêtes
21
+ class ChatRequest(BaseModel):
22
+ message: str
23
+
24
+ @app.post("/chat")
25
+ async def chat(request: ChatRequest):
26
+ inputs = tokenizer(request.message, return_tensors="pt").to("cuda")
27
+ output = model.generate(**inputs, max_length=100)
28
+ response = tokenizer.decode(output[0], skip_special_tokens=True)
29
+ return {"response": response}