n_shot_learning

app.py

@@ -11,9 +11,14 @@ _ = load_dotenv(find_dotenv())
 databricks_token  = os.getenv('DATABRICKS_TOKEN')
 model_uri  = "https://dbc-eb788f31-6c73.cloud.databricks.com/serving-endpoints/Mpt-7b-tester/invocations"
 
-def extract_json(gen_text):
+def extract_json(gen_text, n_shot_learning=0):
     start_index = gen_text.index("### Response:\n{") + 14
-    end_index = gen_text.index("}\n\n### End") + 1
+    if(n_shot_learning > 0) :
+        for i in range(0, n_shot_learning):
+            gen_text = gen_text[start_index:]
+            start_index = gen_text.index("### Response:\n{") + 14
+
+    end_index = gen_text.index("}\n\n### ") + 1
     return gen_text[start_index:end_index]
 
 def score_model(model_uri, databricks_token, prompt):
@@ -43,7 +48,7 @@ def greet(input):
 	response = get_completion(input)
 	gen_text = response["predictions"][0]["generated_text"]
 
-	return extract_json(gen_text)
+	return extract_json(gen_text, 3)
 	#return json.dumps(response)
 
 iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Prompt", lines=3)], outputs="text")