Add FastAPI endpoints for scoring essays and generating prompts

Dockerfile

@@ -0,0 +1,15 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+# FROM python:3.9
+FROM python:latest
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,65 @@
+from transformers import pipeline
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+# Getting the prompt from the prompt.txt file
+prompt_dir = "/prompt.txt"
+prompt = ''
+with open(prompt_dir, 'r') as file:
+    prompt = file.read()
+
+
+def post_process(essay):
+    # Find the index of the first occurrence of the word "Feedback:"
+    feedback_index = essay.find("Feedback:")
+
+    # If "Feedback:" is not found, return the original essay
+    if feedback_index == -1:
+        return essay
+
+    # Find the index of the newline after the first occurrence of "Feedback:"
+    newline_index = essay.find("\n", feedback_index)
+
+    # If newline is not found, return the original essay
+    if newline_index == -1:
+        return essay
+
+    # Return the essay up to the newline after the first occurrence of "Feedback:"
+    return essay[:newline_index]
+
+def pre_process(instruction, essay):
+    text = f"{instruction}\n{essay}"
+    return text
+
+
+pipe = pipeline(
+    "text-generation", 
+    model = "gildead/mistral-aes-414",
+    device_map="auto"
+)
+
+
+class Message(BaseModel):
+    essay: str
+    instruction: str
+        
+app = FastAPI()
+
+
+@app.get("/")
+async def root():
+    return {"message": "Mistral API is running."}
+
+@app.post("/score")
+async def overall(message: Message):
+    text = pre_process(message.instruction, message.essay)
+    result = pipe(
+        f"<s>[INST] {text} [/INST]",
+        max_new_tokens=200, 
+        num_return_sequences=1,)
+    
+    generated_text = result[0]['generated_text']
+    output = generated_text.split('[/INST]', 1)[-1].strip()
+    final_output = post_process(output)
+
+    return {"result": final_output}

main.py

@@ -0,0 +1,71 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import pipeline
+
+class Message(BaseModel):
+    essay: str
+    instruction: str
+
+# Getting the prompt from the prompt.txt file
+prompt_dir = "/prompt.txt"
+prompt = ''
+with open(prompt_dir, 'r') as file:
+    prompt = file.read()
+
+
+def post_process(essay):
+    # Find the index of the first occurrence of the word "Feedback:"
+    feedback_index = essay.find("Feedback:")
+
+    # If "Feedback:" is not found, return the original essay
+    if feedback_index == -1:
+        return essay
+
+    # Find the index of the newline after the first occurrence of "Feedback:"
+    newline_index = essay.find("\n", feedback_index)
+
+    # If newline is not found, return the original essay
+    if newline_index == -1:
+        return essay
+
+    # Return the essay up to the newline after the first occurrence of "Feedback:"
+    return essay[:newline_index]
+
+def pre_process(instruction, essay):
+    text = f"Instruction:{instruction}\nEssay:{essay}"
+    return text
+
+def generate_prompt(input):
+    text = f"""{prompt}\n{input}"""
+    return text
+
+# Initialize your pipeline outside of your route functions
+pipe = pipeline(
+    "text-generation", 
+    model="gildead/mistral-aes-414",
+    device_map="auto"
+)
+
+# Initialize your FastAPI application
+app = FastAPI()
+
+# Define your route functions
+@app.get("/")
+async def root():
+    return {"message": "Mistral API is running."}
+
+@app.post("/score")
+async def overall(message: Message):
+    text = pre_process(message.instruction, message.essay)
+    prompt = generate_prompt(text)
+
+    result = pipe(
+        f"<s>[INST] {prompt} [/INST]",
+        max_new_tokens=200, 
+        num_return_sequences=1,)
+
+    generated_text = result[0]['generated_text']
+    output = generated_text.split('[/INST]', 1)[-1].strip()
+    final_output = post_process(output)
+
+    return {"result": final_output}

prompt.txt

@@ -0,0 +1 @@
+I want you to act as a teacher that grades the essays of the students based on the instructions. The scores range from 1.0 to 5.0 in increments of 0.5. Your task is to give numerical score and text feedback if only needed. Text feedback must be short and one sentence only. Do not be afraid to mark perfect score. You must only reply using this format. \nScore: \nFeedback:

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+bitsandbytes
+accelerate
+fastapi[all] uvicorn