Upload main.py

app/main.py

@@ -12,35 +12,36 @@ from enum import Enum
 from typing import Optional
 
 print("Loading model...")
-llm = Llama(
-      model_path="/models/final-gemma2b_SA-Q8_0.gguf",
+SAllm = Llama(model_path="/models/final-gemma2b_SA-Q8_0.gguf")#,
       # 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
-)
+#)
 
-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
+FIllm = Llama(model_path="/models/final-gemma2b_FI-Q8_0.gguf")
 
-def check_sentiment(text):
-  result = ask(f'Analyze the sentiment of the tweet enclosed in square brackets, determine if it is positive or negative, and return the answer as the corresponding sentiment label "positive" or  "negative" [{text}] =', max_new_tokens=3)
-  return result['choices'][0]['text'].strip()
+# 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 clean_sentiment_response(response_text):
-  result = response_text.strip()
+def check_sentiment(text):
+  prompt = f'Analyze the sentiment of the tweet enclosed in square brackets, determine if it is positive or negative, and return the answer as the corresponding sentiment label "positive" or  "negative" [{text}] ='
+  response = SAllm(prompt, max_tokens=3, stop=["\n"], echo=False, temperature=0.5)
+  # print(response)
+  result = response['choices'][0]['text'].strip()
   if "positive" in result:
     return "positive"
   elif "negative" in result:
     return "negative"
   else:
     return "unknown"
+
   
 print("Testing model...")
 assert "positive" in check_sentiment("ดอกไม้ร้านนี้สวยจัง")
@@ -67,11 +68,15 @@ class SA_Result(str, Enum):
   unknown = "unknown"
 
 class SA_Response(BaseModel):
-  text: Optional[str] = None
   code: int = 200
+  text: Optional[str] = None
   result: SA_Result = None
 
-
+class FI_Response(BaseModel):
+  code: int = 200
+  question: Optional[str] = None
+  answer: str = None
+  config: Optional[dict] = None
 
 @app.get('/')
 def docs():
@@ -88,10 +93,32 @@ def perform_sentiment_analysis(prompt: str = Body(..., embed=True, example="I li
   if prompt:
     try:
       print(f"Checking sentiment for {prompt}")
-      result = clean_sentiment_response(check_sentiment(prompt))
+      result = check_sentiment(prompt)
       print(f"Result: {result}")
       return SA_Response(result=result, text=prompt)
     except Exception as e:
       return HTTPException(500, SA_Response(code=500, result=str(e), text=prompt))
   else:
     return HTTPException(400, SA_Response(code=400, result="Request argument 'prompt' not provided."))
+
+
+@app.post('/FI')
+def ask_gemmaFinanceTH(
+    prompt: str = Body(..., embed=True, example="What's the best way to invest my money"),
+    temperature: float = 0.5, 
+    max_new_tokens: int = 200
+) -> FI_Response:
+  """
+  Ask a finetuned Gemma a finance-related question, just for fun.
+  NOTICE: IT MAY PRODUCE RANDOM/INACCURATE ANSWERS. PLEASE SEEK PROFESSIONAL ADVICE BEFORE DOING ANYTHING SERIOUS.
+  """
+  if prompt:
+    try:
+      print(f'Asking FI with the question "{prompt}"')
+      result = FIllm(prompt, max_tokens=max_new_tokens, temperature=temperature, stop=["###User:", "###Assistant:"], echo=False)
+      print(f"Result: {result}")
+      return FI_Response(answer=result, question=prompt)
+    except Exception as e:
+      return HTTPException(500, FI_Response(code=500, answer=str(e), question=prompt))
+  else:
+    return HTTPException(400, FI_Response(code=400, answer="Request argument 'prompt' not provided."))