Update app.py

app.py

@@ -1,14 +1,18 @@
+
 import streamlit as st
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import time
 import asyncio
 import aiohttp
 import json
+import torch
 
 repo_name = "BeardedMonster/SabiYarn-125M"
-
+device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
 
+model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True).to(device)
+
 # Add sidebar with instructions
 st.sidebar.title("Instructions: How to use")
 st.sidebar.write("""
@@ -48,8 +52,8 @@ st.title("SabiYarn-125M: Generates text in multiple Nigerian languages.")
 
 st.write("**Supported Languages: English, Yoruba, Igbo, Hausa, Pidgin, Efik, Urhobo, Fulfulde, Fulah. \nResults might not be coherent for less represented languages (i.e Efik, \
 Urhobo, Fulfulde, Fulah).**")
-# st.write("**Model is running on CPU RAM. So, token generation might be slower (streaming not enabled).**")
-st.write("**Avg Response time: <3 seconds/ 50 tokens. Response time increase with input length**")
+# st.write("****")
+st.write("**For convenience, you can use chatgpt to copy text and evaluate model output.**")
 st.write("-" * 50)
 
 async def generate_from_api(user_input, generation_config):
@@ -72,13 +76,13 @@ async def generate_from_api(user_input, generation_config):
 # Sample texts
 sample_texts = {
     "Tell me a story in pidgin": "Tell me a story in pidgin",
-    "Oma Ede, Mi ji ogede...": "Oma Ede, Mi ji ogede mi a foroma orhorho edha meji ri eka. "
+    "Oma Ede, Mi ji ogede...": "Oma Ede, Mi ji ogede mi a foroma orhorho edha meji ri eka. ",
     "who are you?": "who are you?",
     "Translate 'how are you?' to Yoruba": "Translate 'how are you?' to Yoruba",
     "Classify the sentiment": "Anyi na-echefu oke ike.",
     "what is the topic of this text": "Africa Free Trade Zone: Kò sí ìdènà láti kó ọjà láti orílẹ̀èdè kan sí òmíràn",
     "diacritize this text: ": "E sun, Alaga, fun ise amalayi ti e n se ni Naijiria. E maa ba a lo, egbon!", 
-    "clean this text": "Di service say sheVbin doE tlykto find out how Zelensky wak^ g3 be before im go visit MykolailXwey flood fKg Junew Zele?sky confam s,y rem don inform  a2 abouo di frrest, im sHy dR_head of SBU bin don update aI abou.Wdi "fiMht againsX traitorsfLZelevsky bin visi# Mykolaiv for June tk see di dam&ge wey di  collapse of KakhoYka dam "ause Rgd den im v:sit again for July afta Russia1heavy bo+bing. Di 8ecurity sBrviceTs1y dem bin know/Mbout di  plan bev9re di ;isit, snd dat onG make put vdditional security measurZs{",
+    "clean this text": "Di service say sheVbin doE tlykto find out how Zelensky wak^ g3 be before im go visit MykolailXwey flood fKg Junew Zele?sky confam s,y rem don inform  a2 abouo di frrest, im sHy dR_head of SBU bin don update aI abou.Wdi 'fiMht againsX traitorsfLZelevsky bin visi# Mykolaiv for June tk see di dam&ge wey di  collapse of KakhoYka dam 'ause Rgd den im v:sit again for July afta Russia1heavy bo+bing. Di 8ecurity sBrviceTs1y dem bin know/Mbout di  plan bev9re di ;isit, snd dat onG make put vdditional security measurZs{",
     "headline of this text": '** Sylvain Itté French ambassador don comot Niger Republic **. Sylvain Itté, di French ambassador for Niger don comot Niamey and currently e dey for flight from Ndjamena to Paris. Sylvain Itté, di French ambassador for Niger don comot Niamey very early dis morning and currently e dey for flight from Ndjamena to Paris.\n\nDi military detain Bazoum and im family for di presidential palace. Niger na former French colony, and France still get 1,500 sojas for di African country.\n\n"France don decide to withdraw dia ambassador. In di next hours our ambassador and several diplomats go return to France," Oga Macron tok.\n\nE add say di military co-operation dey "over" and French troops go leave in "di months to come".\n\n"Dis Sunday we celebrate one new step towards di sovereignty of Niger," di junta tok, for one statement wey AFP news agency quote.\n\nDi decision by Paris dey come afta months of hostility and protest against di presence of French for di kontri, wit regular demonstrations for di capital Niamey.\n\nDi move don scata France operations against Islamist militants for di wider Sahel region and Paris influence for there. But oga Macron tok say "putschists no go hold France hostage,"'
 }
 
@@ -87,8 +91,8 @@ instruction_wrap = {
     "Translate 'how are you?' to Yoruba": "<prompt> Translate 'how are you?' to Yoruba <response>:",
     "who are you?": "<prompt> who are you? <response>:",
     "Classify the sentiment" : "<classify> Anyi na-echefu oke ike. <sentiment>",
-    "clean this text": "<clean>   <pcm>",
-    "diacritize this text: ": "<diacritize>   <yor>"
+    "clean this text": "<clean> Di service say sheVbin doE tlykto find out how Zelensky wak^ g3 be before im go visit MykolailXwey flood fKg Junew Zele?sky confam s,y rem don inform  a2 abouo di frrest, im sHy dR_head of SBU bin don update aI abou.Wdi 'fiMht againsX traitorsfLZelevsky bin visi# Mykolaiv for June tk see di dam&ge wey di  collapse of KakhoYka dam 'ause Rgd den im v:sit again for July afta Russia1heavy bo+bing. Di 8ecurity sBrviceTs1y dem bin know/Mbout di  plan bev9re di ;isit, snd dat onG make put vdditional security measurZs{  <pcm>",
+    "diacritize this text: ": "<diacritize> E sun, Alaga, fun ise amalayi ti e n se ni Naijiria. E maa ba a lo, egbon!  <yor>",
     "what is the topic of this text": "<classify> Africa Free Trade Zone: Kò sí ìdènà láti kó ọjà láti orílẹ̀èdè kan sí òmíràn <topic>",
     'headline of this text': '<title> ** Sylvain Itté French ambassador don comot Niger Republic **. Sylvain Itté, di French ambassador for Niger don comot Niamey and currently e dey for flight from Ndjamena to Paris. Sylvain Itté, di French ambassador for Niger don comot Niamey very early dis morning and currently e dey for flight from Ndjamena to Paris.\n\nDi military detain Bazoum and im family for di presidential palace. Niger na former French colony, and France still get 1,500 sojas for di African country.\n\n"France don decide to withdraw dia ambassador. In di next hours our ambassador and several diplomats go return to France," Oga Macron tok.\n\nE add say di military co-operation dey "over" and French troops go leave in "di months to come".\n\n"Dis Sunday we celebrate one new step towards di sovereignty of Niger," di junta tok, for one statement wey AFP news agency quote.\n\nDi decision by Paris dey come afta months of hostility and protest against di presence of French for di kontri, wit regular demonstrations for di capital Niamey.\n\nDi move don scata France operations against Islamist militants for di wider Sahel region and Paris influence for there. But oga Macron tok say "putschists no go hold France hostage,"  <headline>',
 }
@@ -145,7 +149,17 @@ if st.button("Generate"):
             wrapped_input = wrap_text(user_input)
             print("wrapped_input: ", wrapped_input)
             generation_config["max_new_tokens"]= min(max_new_tokens, 1024 - len(tokenizer.tokenize(wrapped_input)))
-            generated_text = asyncio.run(generate_from_api(wrapped_input, generation_config))
+            try:
+                # Attempt the asynchronous API call
+                generation_config["max_new_tokens"] = min(max_new_tokens, 1024 - len(tokenizer.tokenize(wrapped_input)))
+                generated_text = asyncio.run(generate_from_api(wrapped_input, generation_config))
+            except Exception as e:
+                print(f"API call failed: {e}. Using local model for text generation.")                
+                # Use the locally loaded model for text generation
+                input_ids = tokenizer(wrapped_input, return_tensors="pt")["input_ids"].to(device)
+                output = model.generate(input_ids, **generation_config)
+                generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+                
             full_output = st.empty()
            
             start_time = time.time()
@@ -161,4 +175,4 @@ if st.button("Generate"):
         except Exception as e:
             st.error(f"Error during text generation: {e}")
     else:
-        st.write("Please enter some text to generate.")
+        st.write("Please enter some text to generate.")