| from fastapi import FastAPI, HTTPException, Depends, Request |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| import torch |
| import os |
| import huggingface_hub |
|
|
| app = FastAPI() |
|
|
| EXPECTED_TOKEN = os.environ.get("EXPECTED_TOKEN") |
| HF_TOKEN = os.environ.get('ACCESS_TOKEN') |
|
|
| REPO_ID = "Day23/coder-personal-use" |
| MODEL_FOLDER = "model" |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| huggingface_hub.login(token=HF_TOKEN) |
|
|
| model_dir = huggingface_hub.snapshot_download(repo_id=REPO_ID, allow_patterns=["model/*"]) |
| model_dir = os.path.join(model_dir, 'model') |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True, use_auth_token=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_dir, |
| trust_remote_code=True, |
| device_map=device, |
| ) |
|
|
| @app.post("/generate") |
| async def generate_text(message: str, token: str): |
| """Gera um texto com base na entrada fornecida.""" |
| |
| if not message: |
| raise HTTPException(status_code=400, detail="O campo 'message' é obrigatório.") |
|
|
| if token != EXPECTED_TOKEN: |
| raise HTTPException(status_code=401, detail="Token inválido") |
| |
| messages = [{'role': 'user', 'content': message}] |
| |
| inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(device) |
| |
| with torch.no_grad(): |
| outputs = model.generate( |
| inputs, |
| max_new_tokens=350, |
| do_sample=True, |
| top_k=1, |
| top_p=0.95, |
| num_return_sequences=1, |
| eos_token_id=tokenizer.eos_token_id |
| ) |
| |
| generated_text = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True) |
| |
| return {"response": generated_text} |
