Update main.py

main.py

@@ -6,6 +6,9 @@ from huggingface_hub import InferenceClient
 import random
 from json_repair import repair_json
 import nltk
+import json
+import re
+from word_forms.word_forms import get_word_forms
 
 app = FastAPI()
 
@@ -25,6 +28,10 @@ class InputData(BaseModel):
     json_prompt: str
     history: str = ""
 
+class WordCheckData(BaseModel):
+    string: str
+    word: str
+
 @app.post("/generate-response/")
 async def generate_response(data: InputData) -> Dict[str, Any]:
     client = InferenceClient(model=data.model, token=HF_TOKEN)
@@ -92,4 +99,30 @@ async def generate_response(data: InputData) -> Dict[str, Any]:
 
     except Exception as e:
         print(f"Model {data.model} failed with error: {e}")
-        raise HTTPException(status_code=500, detail=f"Model {data.model} failed to generate response")
+        raise HTTPException(status_code=500, detail=f"Model {data.model} failed to generate response")
+
+@app.post("/check-word/")
+async def check_word(data: WordCheckData) -> Dict[str, Any]:
+    input_string = data.string.lower()
+    word = data.word.lower()
+
+    forms = get_word_forms(word)
+
+    all_forms = set()
+    for words in forms.values():
+        all_forms.update(words)
+
+    # Split the input string into words using regular expression to handle spaces and punctuation
+    words_in_string = re.findall(r'\b\w+\b', input_string)
+
+    found = False
+    for word_in_string in words_in_string:
+        if word_in_string in all_forms:
+            found = True
+            break
+
+    result = {
+        "found": found
+    }
+
+    return result