Feat: AI

app.py

@@ -1,7 +1,46 @@
 from fastapi import FastAPI
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from pydantic import BaseModel
+
+device = "cuda" # the device to load the model onto
+
+model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2-72B-Instruct",
+    torch_dtype="auto",
+    device_map="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-72B-Instruct")
+
+class ChatRequest(BaseModel):
+    prompt: str
 
 app = FastAPI()
 
 @app.get("/")
 def greet_json():
-    return {"Hello": "World!"}
+    return {"Hello": "World!"}
+
+@app.post("/generate_chat")
+def generateAi(request: ChatRequest):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": request.prompt}
+    ]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    model_inputs = tokenizer([text], return_tensors="pt").to(device)
+
+    generated_ids = model.generate(
+        model_inputs.input_ids,
+        max_new_tokens=512
+    )
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    
+    return {"answer": "Hello"}

requirements.txt

@@ -1,2 +1,5 @@
 fastapi
-uvicorn[standard]
+uvicorn[standard]
+transformers
+pydantic
+accelerator