Add FastAPI app and Docker configuration

.dockerignore

@@ -0,0 +1,7 @@
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+.Python
+env/
+venv/

Dockerfile

@@ -0,0 +1,20 @@
+# Use an official Python runtime as a parent image
+FROM python:3.9-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the current directory contents into the container at /app
+COPY . /app
+
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Make port 8000 available to the world outside this container
+EXPOSE 8000
+
+# Define environment variable
+ENV PYTHONUNBUFFERED=1
+
+# Run uvicorn server
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

app/classifier.py

@@ -0,0 +1,18 @@
+from sklearn.linear_model import LogisticRegression
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+def train_classifier(dummy_data):
+    vectorizer = TfidfVectorizer()
+    train_texts, train_labels = zip(*dummy_data)
+    train_vectors = vectorizer.fit_transform(train_texts)
+    classifier = LogisticRegression()
+    classifier.fit(train_vectors, train_labels)
+    return classifier, vectorizer
+
+def classify_text(text: str, classifier, vectorizer) -> str:
+    try:
+        transformed_data = vectorizer.transform([text])
+        category = classifier.predict(transformed_data)[0]
+        return category
+    except Exception as e:
+        return str(e)

app/main.py

@@ -0,0 +1,47 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from app.model_utils import load_model_and_tokenizer, generate_summary
+from app.classifier import train_classifier, classify_text
+
+app = FastAPI()
+
+# Load model and tokenizer for the /rag endpoint
+model_name = "sshleifer/distilbart-cnn-6-6"  # Example model 
+model, tokenizer = load_model_and_tokenizer(model_name)
+
+# Dummy data and classifier for the /classification endpoint
+dummy_data = [
+    ("I feel very sad and hopeless.", "Depression"),
+    ("I have trouble sleeping at night.", "Insomnia"),
+    ("I am constantly worrying about everything.", "Anxiety"),
+    ("I feel energetic and happy.", "Happiness"),
+    ("My mood swings a lot and I feel irritable.", "Mood Disorder")
+]
+
+classifier, vectorizer = train_classifier(dummy_data)
+
+class Prompt(BaseModel):
+    prompt: str
+
+class ClassificationInput(BaseModel):
+    data: str
+
+@app.post("/rag")
+def rag_endpoint(prompt: Prompt):
+    try:
+        summary = generate_summary(prompt.prompt, model, tokenizer)
+        return {"summary": summary}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/classification")
+def classification_endpoint(input: ClassificationInput):
+    try:
+        category = classify_text(input.data, classifier, vectorizer)
+        return {"category": category}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

app/model_utils.py

@@ -0,0 +1,17 @@
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+def load_model_and_tokenizer(model_name: str):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return model, tokenizer
+
+def generate_summary(prompt: str, model, tokenizer) -> str:
+    context_prompt = f"Provide a brief, informative article addressing the following mental health concern: {prompt}"
+    
+    try:
+        inputs = tokenizer(context_prompt, return_tensors="pt", truncation=True, padding=True)
+        summary_ids = model.generate(inputs["input_ids"], max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
+        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        return summary
+    except Exception as e:
+        return str(e)

requirements.txt

@@ -0,0 +1,5 @@
+fastapi
+uvicorn
+transformers
+torch
+scikit-learn