kk

app.py

@@ -83,7 +83,13 @@ def lang_model_response(prompt):
   print(f"output is : {output}") 
   output_tmp = output[0]['generated_text']
   print(f"output_tmp is: {output_tmp}")
-  solution = output_tmp.split("\n")[2]   
+  if '?' in prompt:
+    solution = output_tmp.split("?")[1]
+  elif '.' in prompt:
+    solution = output_tmp.split(".")[1]
+  else:
+    solution = output_tmp.split("\n")[1]
+    
   print(f"Final response after splits is: {solution}")
   return solution