Change data provider

app.py

@@ -1,34 +1,29 @@
-from typing import List, Dict, Any
-import zipfile
-import os
-import warnings
-from openai import OpenAI
+import os, warnings
 from dotenv import load_dotenv
+from schemas import *
+
+os.environ["CURL_CA_BUNDLE"] = ""
+warnings.filterwarnings("ignore")
+load_dotenv()
+
+from datasets import load_dataset
 import bm25s
-from fastapi.staticfiles import StaticFiles
-from nltk.stem import WordNetLemmatizer
-import nltk
+from bm25s.hf import BM25HF
+
 from fastapi import FastAPI
-from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
-import numpy as np
-from pydantic import BaseModel
-from sklearn.preprocessing import MinMaxScaler
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
 
-load_dotenv()
+from sklearn.preprocessing import MinMaxScaler
+import numpy as np
 
-nltk.download('wordnet')
-if os.path.exists("bm25s.zip"):
-    with zipfile.ZipFile("bm25s.zip", 'r') as zip_ref:
-        zip_ref.extractall(".")
-    bm25_engine = bm25s.BM25.load("3gpp_bm25_docs", load_corpus=True)
-lemmatizer = WordNetLemmatizer()
-llm = OpenAI(api_key=os.environ.get("GEMINI"), base_url="https://generativelanguage.googleapis.com/v1beta/openai/")
+import litellm
 
-warnings.filterwarnings("ignore")
+bm25_index = BM25HF.load_from_hub("OrganizedProgrammers/3GPPBM25IndexSections", load_corpus=True, token=os.environ["HF_TOKEN"])
 
 app = FastAPI(title="RAGnarok", 
-              description="API to search specifications for RAG")
+              description="Speak with the specifications")
 
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
@@ -44,46 +39,22 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-class SearchRequest(BaseModel):
-    keyword: str
-    threshold: int
-
-class SearchResponse(BaseModel):
-    results: List[Dict[str, Any]]
-
-class ChatRequest(BaseModel):
-    messages: List[Dict[str, str]]
-    model: str
-
-class ChatResponse(BaseModel):
-    response: str
-
 @app.get("/")
-async def main_menu():
+def main_menu():
     return FileResponse(os.path.join("templates", "index.html"))
 
-@app.post("/chat", response_model=ChatResponse)
-def question_the_sources(req: ChatRequest):
-    model = req.model
-    resp = llm.chat.completions.create(
-        messages=req.messages,
-        model=model
-    )
-    return ChatResponse(response=resp.choices[0].message.content)
-
 @app.post("/search", response_model=SearchResponse)
 def search_specifications(req: SearchRequest):
     keywords = req.keyword
     threshold = req.threshold
-    query = lemmatizer.lemmatize(keywords)
     results_out = []
-    query_tokens = bm25s.tokenize(query)
-    results, scores = bm25_engine.retrieve(query_tokens, k=len(bm25_engine.corpus))
+    query_tokens = bm25s.tokenize(keywords)
+    results, scores = bm25_index.retrieve(query_tokens, k=len(bm25_index.corpus))
 
     def calculate_boosted_score(metadata, score, query):
-        title = {lemmatizer.lemmatize(metadata['title']).lower()}
-        q = {query.lower()}
-        spec_id_presence = 0.5 if len(q & {metadata['id']}) > 0 else 0
+        title = set(metadata['title'].lower().split())
+        q = set(query.lower().split())
+        spec_id_presence = 0.5 if metadata['id'].lower() in q else 0
         booster = len(q & title) * 0.5
         return score + spec_id_presence + booster
 
@@ -96,7 +67,7 @@ def search_specifications(req: SearchRequest):
         score = scores[0, i]
         spec = doc["metadata"]["id"]
 
-        boosted_score = calculate_boosted_score(doc['metadata'], score, query)
+        boosted_score = calculate_boosted_score(doc['metadata'], score, keywords)
 
         if spec not in spec_scores or boosted_score > spec_scores[spec]:
             spec_scores[spec] = boosted_score
@@ -135,4 +106,14 @@ def search_specifications(req: SearchRequest):
             break
         results_out.append({'id': metadata['id'], 'title': metadata['title'], 'section': metadata['section_title'], 'content': details['doc']['text'], 'similarity': int(details['normalized_score']*100)})
     
-    return SearchResponse(results=results_out)
+    return SearchResponse(results=results_out)
+
+@app.post("/chat", response_model=ChatResponse)
+def questions_the_sources(req: ChatRequest):
+    model = req.model
+    resp = litellm.completion(
+        model=f"gemini/{model}",
+        messages=req.messages,
+        api_key=os.environ["GEMINI"]
+    )
+    return ChatResponse(response=resp.choices[0].message.content)

requirements.txt

@@ -1,8 +1,8 @@
-openai
 fastapi
 uvicorn[standard]
 python-dotenv
 bm25s[full]
-nltk
+scikit-learn
+litellm
 numpy
-scikit-learn
+datasets

schemas.py

@@ -0,0 +1,16 @@
+from typing import *
+from pydantic import BaseModel
+
+class SearchRequest(BaseModel):
+    keyword: str
+    threshold: int
+
+class SearchResponse(BaseModel):
+    results: List[Dict[str, Any]]
+
+class ChatRequest(BaseModel):
+    messages: List[Dict[str, str]]
+    model: str
+
+class ChatResponse(BaseModel):
+    response: str