Spaces:
Runtime error
Runtime error
File size: 1,447 Bytes
8d6c4b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModel
import uvicorn
import gradio as gr
# Initialize FastAPI app
app = FastAPI()
# Load pre-trained model and tokenizer
model_name = "bert-base-uncased" # You can change this to another model
try:
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
except Exception as e:
print(f"Error loading model: {e}")
class TextRequest(BaseModel):
text: str
# Function to generate embeddings
def get_embeddings(text: str):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
outputs = model(**inputs)
embeddings = outputs.last_hidden_state.mean(dim=1) # Pooling
return embeddings.numpy().tolist() # Convert to list for API response
@app.post("/get-embedding/")
async def get_embedding(request: TextRequest):
text = request.text
embeddings = get_embeddings(text)
return {"embedding": embeddings}
def gradio_interface(text):
return get_embeddings(text)
grn = gr.Interface(fn=gradio_interface, inputs="text", outputs="json", title="Text Embedding Generator")
grn.launch(server_name="0.0.0.0", server_port=7860, share=True)
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860) # Port 7860 for Hugging Face Spaces
|