Upload main.py

app/main.py

@@ -10,25 +10,16 @@ from pydantic import BaseModel
 from enum import Enum
 from typing import Optional
 
+# MODEL LOADING, FUNCTIONS, AND TESTING
+
 print("Loading model...")
 SAllm = Llama(model_path="/models/final-gemma2b_SA-Q8_0.gguf", mmap=False, mlock=True)
+FIllm = Llama(model_path="/models/final-gemma7b_FI-Q8_0.gguf", mmap=False, mlock=True)
       # n_gpu_layers=28, # Uncomment to use GPU acceleration
       # seed=1337, # Uncomment to set a specific seed
       # n_ctx=2048, # Uncomment to increase the context window
 #)
 
-FIllm = Llama(model_path="/models/final-gemma7b_FI-Q8_0.gguf", mmap=False, mlock=True)
-
-# def ask(question, max_new_tokens=200):
-#   output = llm(
-#     question, # Prompt
-#     max_tokens=max_new_tokens, # Generate up to 32 tokens, set to None to generate up to the end of the context window
-#     stop=["\n"], # Stop generating just before the model would generate a new question
-#     echo=False, # Echo the prompt back in the output
-#     temperature=0.0,
-#   )
-#   return output
-
 def extract_restext(response):
   return response['choices'][0]['text'].strip()
 
@@ -49,15 +40,17 @@ def check_sentiment(text):
   else:
     return "unknown"
 
-  
+# TESTING THE MODEL
 print("Testing model...")
 assert "positive" in check_sentiment("ดอกไม้ร้านนี้สวยจัง")
 assert ask_fi("Hello!, How are you today?")
 print("Ready.")
 
+
+# START OF FASTAPI APP
 app = FastAPI(
-    title = "GemmaSA_2b",
-    description="A simple sentiment analysis API for the Thai language, powered by a finetuned version of Gemma-2b",
+    title = "Gemma Finetuned API",
+    description="Gemma Finetuned API for Sentiment Analysis and Finance Questions.",
     version="1.0.0",
 )
 
@@ -70,6 +63,8 @@ app.add_middleware(
     allow_headers=["*"]
 )
 
+
+# API DATA CLASSES
 class SA_Result(str, Enum):
   positive = "positive"
   negative = "negative"
@@ -86,17 +81,15 @@ class FI_Response(BaseModel):
   answer: str = None
   config: Optional[dict] = None
 
+
+# API ROUTES
 @app.get('/')
 def docs():
   "Redirects the user from the main page to the docs."
   return responses.RedirectResponse('./docs')
 
-@app.get('/add/{a}/{b}')
-def add(a: int,b: int):
-    return a + b
-
-@app.post('/SA')
-def perform_sentiment_analysis(prompt: str = Body(..., embed=True, example="I like eating fried chicken")) -> SA_Response:
+@app.post('/classifications/sentiment')
+async def perform_sentiment_analysis(prompt: str = Body(..., embed=True, example="I like eating fried chicken")) -> SA_Response:
   """Performs a sentiment analysis using a finetuned version of Gemma-7b"""
   if prompt:
     try:
@@ -110,8 +103,8 @@ def perform_sentiment_analysis(prompt: str = Body(..., embed=True, example="I li
     return HTTPException(400, SA_Response(code=400, result="Request argument 'prompt' not provided."))
 
 
-@app.post('/FI')
-def ask_gemmaFinanceTH(
+@app.post('/questions/finance')
+async def ask_gemmaFinanceTH(
     prompt: str = Body(..., embed=True, example="What's the best way to invest my money"),
     temperature: float = Body(0.5, embed=True), 
     max_new_tokens: int = Body(200, embed=True)