Update app.py

app.py

@@ -47,8 +47,25 @@ def print_output(input):
     # outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True)
     # tokenizer.batch_decode(outputs)
     
-    outputs = model.generate(**inputs, max_new_tokens = 100, use_cache = True)
-    return tokenizer.batch_decode(outputs)
+    # outputs = model.generate(**inputs, max_new_tokens = 100, use_cache = True)
+    # return tokenizer.batch_decode(outputs)
+
+
+    outputs = model.generate(**inputs, max_new_tokens = 200, use_cache = True, temperature=0.5)
+
+    n1 = tokenizer.batch_decode(outputs)[0].index('### Response:')
+    output_ = tokenizer.batch_decode(outputs)[0][n1+13:].replace('.com','')
+    n2 = output_.index('\n\n')
+
+    dom_names = []
+    
+    for name in output_[:n2].split('\n'):
+      if name != '':
+          dom_names.append(name.split('.')[1])
+    
+    s = '\n'.join(dom_names)
+
+    return s
 
 
 # Streamlit app UI