work on dialogue models

README.md

@@ -1,6 +1,6 @@
 ---
-title: Llama 2 13b Chat
-emoji: 🦙
+title: Chat with Masakhane Dialogue Models
+emoji: 🌍
 colorFrom: indigo
 colorTo: pink
 sdk: gradio
@@ -8,12 +8,7 @@ sdk_version: 3.37.0
 app_file: app.py
 pinned: false
 license: other
-suggested_hardware: a10g-small
 duplicated_from: huggingface-projects/llama-2-13b-chat
 ---
 
-# LLAMA v2 Models
-
-Llama v2 was introduced in [this paper](https://arxiv.org/abs/2307.09288).
-
-This Space demonstrates [Llama-2-13b-chat-hf](meta-llama/Llama-2-13b-chat-hf) from Meta. Please, check the original model card for details.
+# Chat with Masakhane Dialogue Models

app.py

@@ -12,7 +12,9 @@ MAX_INPUT_TOKEN_LENGTH = 4000
 
 DESCRIPTION = """
 # Masakhane Dialogue Models
-This Space demonstrates the dialogue models for Nigerian Pidgin, an African langage.\
+
+This Space demonstrates the dialogue models for Nigerian Pidgin, an African langage.\n
+
 🔎 For more about visit [our homepage](https://www.masakhane.io/).
 
 """
@@ -21,7 +23,7 @@ This Space demonstrates the dialogue models for Nigerian Pidgin, an African lang
 
 
 if not torch.cuda.is_available():
-    DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'
+    DESCRIPTION += '\n<p>Running on CPU 🥶 This demo will be very slow on CPU.</p>'
 
 
 def clear_and_save_textbox(message: str) -> tuple[str, str]:

model.py

@@ -5,18 +5,18 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIt
 
 model_id = 'tosin/dialogpt_afriwoz_pidgin'
 
-if torch.cuda.is_available():
-    config = AutoConfig.from_pretrained(model_id)
-    config.pretraining_tp = 1
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        config=config,
-        torch_dtype=torch.float16,
-        load_in_4bit=True,
-        device_map='auto'
+#if torch.cuda.is_available():
+config = AutoConfig.from_pretrained(model_id)
+config.pretraining_tp = 1
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    config=config,
+    #torch_dtype=torch.float16,
+    #load_in_4bit=True,
+    device_map='cpu'
     )
-else:
-    model = None
+#else:
+#    model = None
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
@@ -51,10 +51,11 @@ def run(message: str,
         top_p: float = 0.95,
         top_k: int = 50) -> Iterator[str]:
     prompt = get_prompt(message, chat_history, system_prompt)
-    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')
+    #inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')
+    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False)
 
     streamer = TextIteratorStreamer(tokenizer,
-                                    timeout=10.,
+                                    timeout=40.,
                                     skip_prompt=True,
                                     skip_special_tokens=True)
     generate_kwargs = dict(
@@ -62,8 +63,8 @@ def run(message: str,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
+        #top_p=top_p,
+        #top_k=top_k,
         temperature=temperature,
         num_beams=1,
     )