Upload main.py

app/main.py

@@ -30,11 +30,14 @@ FIllm = Llama(model_path="/models/final-gemma2b_FI-Q8_0.gguf")
 #   )
 #   return output
 
+def extract_restext(response):
+  return response['choices'][0]['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()
+  result = extract_restext(response)
   if "positive" in result:
     return "positive"
   elif "negative" in result:
@@ -115,7 +118,7 @@ def ask_gemmaFinanceTH(
   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)
+      result = extract_restext(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: