Update inference_server.py

inference_server.py

@@ -1,22 +1,30 @@
+import pickle
 import logging 
-from sklearn.linear_model import SGDClassifier
+from sklearn.feature_extraction.text import TfidVectorizer
+from sklearn.pipeline import Pipeline
+from sklearn.native_bayes import MultinomialNB
 import uvicorn
 from fastapi import FastAPI
 
 app = FastAPI()
 
+strings = set() # Set to store all input strings
+
 def predict(input_text: str):
-    data = [[ord(c) for c in input_text]] # Convert the string to a list of ASCII values
-    model = train(data)
-    # Make a prediction
-    prediction = model.predict([[ord(c) for c in 'abc']]) # Convert the input string to a list of ASCII values
-    return {"prediction": prediction}
+    # Add the new input string to the set of strings
+    strings.add(input_text)
+    # Train a new model using all strings in the set
+    model = Pipeline([
+        ('vectorizer', TfidVectorizer()),
+        ('classifier', MultinomialNB())
+    ])
+    model.fit(list(strings), list(strings))
 
-def train(X):
-    model = SGDClassifier()
-    model.fit(X, X) # In this case, we are using the input data as the labels
-    return model
+    # Make a prediction on the new input string
+    prediction = model.predict([input_text])[0]
 
+    return {"prediction": prediction}
+    
 # Here you can do things such as load your models
 
 @app.get("/")
@@ -25,7 +33,7 @@ def read_root(input_text):
     try:
         result = predict(input_text)
         logging.info("Prediction made: %s", result)
-        return {"result": 1}
+        return result
     except Exception as e:
         logging.error("An error occured: %s", e)
         return {"error": str(e)}