30 lines
1011 B
Python
30 lines
1011 B
Python
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
from transformers import AutoTokenizer, AutoModel
|
|
import torch
|
|
import uvicorn # <-- Добавили импорт
|
|
|
|
app = FastAPI(title="GraphCodeBERT Vectorizer")
|
|
|
|
# Загружаем модель глобально при старте приложения
|
|
model_name = "microsoft/graphcodebert-base"
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
model = AutoModel.from_pretrained(model_name)
|
|
|
|
class ChunkRequest(BaseModel):
|
|
text: str
|
|
|
|
@app.post("/vectorize")
|
|
def vectorize(request: ChunkRequest):
|
|
inputs = tokenizer(request.text, return_tensors="pt", truncation=True, max_length=512)
|
|
|
|
with torch.no_grad():
|
|
outputs = model(**inputs)
|
|
|
|
vector = outputs.last_hidden_state[:, 0, :].squeeze().tolist()
|
|
return {"vector": vector}
|
|
|
|
# <-- Добавили блок запуска
|
|
if __name__ == "__main__":
|
|
print("Запускаем сервер на порту 8000...")
|
|
uvicorn.run(app, host="0.0.0.0", port=8000) |