app.py

@@ -1,6 +1,18 @@
+"""Try out gradio.Chatinterface.
+
+colab gradio-chatinterface.
+
+%%writefile reuirements.txt
+gradio
+transformers
+sentencepiece
+torch
+
+"""
+# pylint: disable=line-too-long, missing-module-docstring, missing-function-docstring
 # import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
+from transformers import AutoModel, AutoTokenizer  # AutoModelForCausalLM,
 
 # device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -12,7 +24,7 @@ tokenizer = AutoTokenizer.from_pretrained(
     "THUDM/chatglm2-6b-int4", trust_remote_code=True
 )
 chat_model = AutoModel.from_pretrained(
-    "THUDM/chatglm2-6b-int4", trust_remote_code=True
+    "THUDM/chatglm2-6b-int4", trust_remote_code=True  # 3.92G
 ).float()
 
 
@@ -21,7 +33,7 @@ def chat(message, history):
     # inputs = tokenizer(prompt, return_tensors="pt").to(device=device)
     # output = model.generate(**inputs, do_sample=True, top_p=0.95, top_k=0, max_new_tokens=256)
     # return tokenizer.decode(output[0], skip_special_tokens=True)
-    for response, history in chat_model.stream_chat(
+    for response, _ in chat_model.stream_chat(
         tokenizer, message, history, max_length=2048, top_p=0.7, temperature=0.95
     ):
         yield response
@@ -34,5 +46,5 @@ gr.ChatInterface(
     examples=[
         ["test me"],
     ],
-    theme=gr.themes.Soft(),
+    theme=gr.themes.Glass(text_size="sm", spacing_size="sm"),
 ).queue(max_size=2).launch()