Update functionality.py

functionality.py

@@ -112,9 +112,9 @@ def get_collection() -> chromadb.Collection:
         papers = pd.read_csv("hf://datasets/somosnlp-hackathon-2022/scientific_papers_en/scientific_paper_en.csv")
         logging.info(f"The data downloaded from url.")
         papers = papers.drop(['id'], axis=1)
-        papers = papers.iloc[:100]
+        papers = papers.iloc[:200]
 
-        for i in range(100):
+        for i in range(200):
             paper = papers.iloc[i]
             idx = paper.name
 
@@ -231,7 +231,7 @@ def _chat_stream(initial_text: str, parts: list):
         )
         inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
         streamer = TextIteratorStreamer(
-            tokenizer=tokenizer, skip_prompt=True, timeout=160.0, skip_special_tokens=True
+            tokenizer=tokenizer, skip_prompt=True, timeout=60.0, skip_special_tokens=True
         )
         generation_kwargs = {
             **inputs,
@@ -253,44 +253,22 @@ def predict(goal: str, parts: list, context: str):
             logging.info("No context was provided!")
         elif goal == 'Fix Academic Style':
             formal_text = ""
-            try:
-                for new_text in fix_academic_style(context):
-                    formal_text = new_text
-                    yield formal_text
-                if not formal_text:
-                    yield "Generation failed or timed out. Please try again!"
-
-                logging.info(f"\n---Academic style corrected:---\n {formal_text}\n")
-            except Exception as e:
-                logging.error(f"Error in 'Fix Academic Style' occured: {e}")
-                yield "Try to wait a little bit and resend your request!"
+            for new_text in fix_academic_style(context):
+                formal_text = new_text
+                yield formal_text
 
+            logging.info(f"\n---Academic style corrected:---\n {formal_text}\n")
         elif goal == 'Fix Grammar':
-            try:
-                full_response = ""
-                for new_text in fix_grammar(context):
-                    full_response = new_text
-                    yield full_response
-                
-                if not full_response:
-                    yield "Generation failed or timed out. Please try again!"
-                
-                logging.info(f"\n---Grammar corrected:---\n{full_response}\n")
-            except Exception as e:
-                logging.error(f"Error in 'Fix Grammar' occured: {e}")
-                yield "Try to wait a little bit and resend your request!"
-
+            full_response = ""
+            for new_text in fix_grammar(context):
+                full_response = new_text
+                yield full_response
+            
+            logging.info(f"\n---Grammar corrected:---\n{full_response}\n")
         else:
-            try:
-                full_response = ""
-                for new_text in _chat_stream(context, parts):
-                    full_response = new_text
-                    yield full_response
-                
-                if not full_response:
-                    yield "Generation failed or timed out. Please try again!"
+            full_response = ""
+            for new_text in _chat_stream(context, parts):
+                full_response = new_text
+                yield full_response
 
-                logging.info(f"\nThe text was generated!\n{full_response}")
-            except Exception as e:
-                logging.error(f"Error in 'Write Text' occured: {e}")
-                yield "Try to wait a little bit and resend your request!"
+            logging.info(f"\nThe text was generated!\n{full_response}")