test

app.py

@@ -50,8 +50,8 @@ for i, j in zip(ents, ents_prompt):
     print(i, j)
 
 model_mapping = {
-    'gpt3.5': 'gpt2',
-    #'vicuna-7b': 'lmsys/vicuna-7b-v1.3',
+    #'gpt3.5': 'gpt2',
+    'vicuna-7b': 'lmsys/vicuna-7b-v1.3',
     #'vicuna-13b': 'lmsys/vicuna-13b-v1.3',
     #'vicuna-33b': 'lmsys/vicuna-33b-v1.3',
     #'fastchat-t5': 'lmsys/fastchat-t5-3b-v1.0',
@@ -94,8 +94,8 @@ with open('demonstration_3_42_parse.txt', 'r') as f:
 theme = gr.themes.Soft()
 
 
-gpt_pipeline = pipeline(task="text-generation", model="gpt2")
-#vicuna7b_pipeline = pipeline(task="text2text-generation", model="lmsys/vicuna-7b-v1.3")
+#gpt_pipeline = pipeline(task="text-generation", model="gpt2")
+vicuna7b_pipeline = pipeline(task="text2text-generation", model="lmsys/vicuna-7b-v1.3")
 #vicuna13b_pipeline = pipeline(task="text2text-generation", model="lmsys/vicuna-13b-v1.3")
 #vicuna33b_pipeline = pipeline(task="text2text-generation", model="lmsys/vicuna-33b-v1.3")
 #fastchatT5_pipeline = pipeline(task="text2text-generation", model="lmsys/fastchat-t5-3b-v1.0")
@@ -117,33 +117,33 @@ def process_text(model_name, task, text):
     for gid in tqdm(gid_list, desc='Query'):
             text = ptb[gid]['text']
     
-    if model_name == 'gpt3.5': 
+    if model_name == 'vicuna-7b': 
         if task == 'POS':
             strategy1 = template_all.format(text)
             strategy2 = prompt2_pos.format(text)
             strategy3 = demon_pos
 
-            response1 = gpt_pipeline(strategy1)[0]['generated_text']
-            response2 = gpt_pipeline(strategy2)[0]['generated_text']
-            response3 = gpt_pipeline(strategy3)[0]['generated_text']
+            response1 = vicuna7b_pipeline(strategy1)[0]['generated_text']
+            response2 = vicuna7b_pipeline(strategy2)[0]['generated_text']
+            response3 = vicuna7b_pipeline(strategy3)[0]['generated_text']
             return (response1, response2, response3)
         elif task == 'Chunking':
             strategy1 = template_all.format(text)
             strategy2 = prompt2_chunk.format(text)
             strategy3 = demon_chunk
 
-            response1 = gpt_pipeline(strategy1)[0]['generated_text']
-            response2 = gpt_pipeline(strategy2)[0]['generated_text']
-            response3 = gpt_pipeline(strategy3)[0]['generated_text']
+            response1 = vicuna7b_pipeline(strategy1)[0]['generated_text']
+            response2 = vicuna7b_pipeline(strategy2)[0]['generated_text']
+            response3 = vicuna7b_pipeline(strategy3)[0]['generated_text']
             return (response1, response2, response3)
         elif task == 'Parsing':
             strategy1 = template_all.format(text)
             strategy2 = prompt2_parse.format(text)
             strategy3 = demon_parse
             
-            response1 = gpt_pipeline(strategy1)[0]['generated_text']
-            response2 = gpt_pipeline(strategy2)[0]['generated_text']
-            response3 = gpt_pipeline(strategy3)[0]['generated_text']
+            response1 = vicuna7b_pipeline(strategy1)[0]['generated_text']
+            response2 = vicuna7b_pipeline(strategy2)[0]['generated_text']
+            response3 = vicuna7b_pipeline(strategy3)[0]['generated_text']
             return (response1, response2, response3)
 
 # Gradio interface