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)