semanticscore / main.py
saicharan1234's picture
Update main.py
a3f0405
import os
# cwd = os.getcwd()
# cachedir = cwd+'/cache'
# os.mkdir(cachedir)
# os.environ['TRANSFORMERS_CACHE'] = "/app/yo/cache"
from fastapi import FastAPI, HTTPException
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer, util
app = FastAPI()
# Load the SentenceTransformer model
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
class StringInput(BaseModel):
sentence1: str
sentence2: str
class ListInput(BaseModel):
sentences1: list
sentences2: list
class MatrixInput(BaseModel):
matrix: list
# Configure CORS settings
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # You should restrict this to specific origins in production
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.post("/calculate-cosine-similarity-tabulated")
async def calculate_cosine_similarity_tabulated(data: ListInput):
try:
# Encode the sentences
embeddings1 = model.encode(data.sentences1)
embeddings2 = model.encode(data.sentences2)
# Calculate cosine similarity scores
similarity_table = []
# Add the first row with the specified headers
first_row = [""] + ["PO1", "PO2", "PO3", "PO4", "PO5", "PO6", "PO7", "PO8", "PO9", "PO10", "PO11", "PO12", "PSO1", "PSO2", "PSO3"]
similarity_table.append(first_row)
# Iterate through the sentences and calculate similarity scores
for i, sentence1 in enumerate(embeddings1):
row = ["CO" + str(i + 1)] # Start the row with "S1", "S2", ...
for sentence2 in embeddings2:
score = util.pytorch_cos_sim(sentence1, sentence2).item()
score = max(score, 0)
formatted_score = f"{score:.3f}"
row.append(formatted_score)
# Add the row to the similarity table
similarity_table.append(row)
return {"cosine_similarity_table": similarity_table}
except Exception as e:
# Handle any exceptions that might occur during calculation
raise HTTPException(status_code=500, detail="An error occurred during calculation")
@app.post("/calculate-cosine-similarity")
async def calculate_cosine_similarity(data: StringInput):
try:
# Encode the sentences
my_embedding = model.encode(data.sentence1)
embeddings = model.encode(data.sentence2)
# Calculate cosine similarity
cos_sim = util.pytorch_cos_sim(my_embedding, embeddings).item()
return {"cosine_similarity": cos_sim}
except Exception as e:
raise HTTPException(status_code=500, detail="An error occurred during calculation")