|
|
from fastapi import FastAPI |
|
|
from pydantic import BaseModel |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
import torch |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
|
|
|
MODEL_NAME = "TheBloke/vicuna-7B-1.1-HF" |
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) |
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", torch_dtype=torch.float16) |
|
|
|
|
|
|
|
|
chat_histories = {} |
|
|
|
|
|
class ChatRequest(BaseModel): |
|
|
user_id: str |
|
|
message: str |
|
|
|
|
|
@app.post("/chat") |
|
|
def chat_endpoint(request: ChatRequest): |
|
|
user_id = request.user_id |
|
|
message = request.message |
|
|
|
|
|
if user_id not in chat_histories: |
|
|
chat_histories[user_id] = [] |
|
|
|
|
|
conversation = "NOVA AI: أنا كوميدي ومغربي. نفهم أي حاجة.\n" |
|
|
for q, a in chat_histories[user_id]: |
|
|
conversation += f"User: {q}\nNOVA AI: {a}\n" |
|
|
conversation += f"User: {message}\nNOVA AI:" |
|
|
|
|
|
inputs = tokenizer(conversation, return_tensors="pt").to(model.device) |
|
|
outputs = model.generate(**inputs, max_new_tokens=200) |
|
|
response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("NOVA AI:")[-1].strip() |
|
|
|
|
|
chat_histories[user_id].append((message, response)) |
|
|
if len(chat_histories[user_id]) > 10: |
|
|
chat_histories[user_id] = chat_histories[user_id][-10:] |
|
|
|
|
|
return {"response": response} |
|
|
|
