Spaces:

moraxgiga
/

rag-demo

Runtime error

rag-demo / main.py

Update main.py

5c5054d verified 5 months ago

No virus

1.2 kB

	from fastapi import FastAPI, HTTPException
	from pydantic import BaseModel, root_validator
	from transformers import AutoModel
	from typing import List
	import os, platform, time
	from fastapi.middleware.cors import CORSMiddleware # Import CORSMiddleware

	# Your model loading and app initialization code
	model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
	app = FastAPI()

	# Set up CORS middleware configuration
	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"], # Allows all origins
	allow_credentials=True,
	allow_methods=["*"], # Allows all methods
	allow_headers=["*"], # Allows all headers
	)

	# Your Validation model and endpoint definitions remain the same
	class Validation(BaseModel):
	prompt: List[str]

	@app.post("/jina_embedding")
	async def generate_embeddings(item: Validation):
	start_time = time.time()
	embeddings = model.encode(item.prompt).tolist()
	end_time = time.time()
	time_taken = end_time - start_time # Calculate the time taken

	return {
	"embeddings": embeddings,
	"time_taken": f"{time_taken:.2f} seconds",
	"Number_of_sentence_processed": len(item.prompt)
	}