Добавьте файлы проекта.
This commit is contained in:
30
PythonProject/main.py
Normal file
30
PythonProject/main.py
Normal file
@@ -0,0 +1,30 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user