Update app.py

app.py

@@ -1,69 +1,24 @@
-import os
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-model_name = "meta-llama/Llama-2-13b-chat-hf"  # 将此处替换为你想使用的模型的名称
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-class Conversation:
-    def __init__(self, prompt, round):
-        self.prompt = prompt
-        self.round = round
-        self.messages = []
-        self.messages.append({"role": "system", "content": self.prompt})
-
-    def ask(self, question):
-        try:
-            self.messages.append({"role": "user", "content": question})
-            input_text = self._build_message(self.messages)
-            encoded_input = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')
-            output = model.generate(encoded_input, max_length=200, temperature=0.5)
-            message = tokenizer.decode(output[:, encoded_input.shape[-1]:][0], skip_special_tokens=True)
-        except Exception as e:
-            print(e)
-            return e
-
-        self.messages.append({"role": "assistant", "content": message})
-
-        if len(self.messages) > self.round * 2 + 1:
-            text = self._build_message(self.messages)
-            self.messages = []
-            self.messages.append({"role": "system", "content": text})
-        return message
-
-    def _build_message(self, messages):
-        text = ""
-        for message in messages:
-            if message["role"] == "user":
-                text += "User : " + message["content"] + "\n\n"
-            if message["role"] == "assistant":
-                text += "Assistant : " + message["content"] + "\n\n" 
-        return text
-
-prompt = """你是一个大数据和AI领域的专家，用中文回答大数据和AI的相关问题。你的回答需要满足以下要求:
-1. 你的回答必须是中文
-2. 回答限制在200个字以内
-3. 拒绝回答违反社会道德和法律的问题"""
-
-conv = Conversation(prompt, 3)
-
-def answer(question, history=[]):
-    history.append(question)
-    message = conv.ask(question)
-    history.append(message)
-    responses = [(u,b) for u,b in zip(history[::2], history[1::2])]
-    print(responses)
-    return responses, history
-
-with gr.Blocks(css="#chatbot{height:300px} .overflow-y-auto{height:500px}") as rxbot:
-    chatbot = gr.Chatbot(elem_id="chatbot")
-    state = gr.State([])
-
-    with gr.Row():
-        txt = gr.Textbox(show_label=False, placeholder="请输入你的问题").style(container=False)
-
-    txt.submit(answer, [txt, state], [chatbot, state])
-    
-
-rxbot.launch()
+from transformers import AutoTokenizer
+import transformers
+import torch
+
+model = "meta-llama/Llama-2-7b-chat-hf"
+
+tokenizer = AutoTokenizer.from_pretrained(model)
+pipeline = transformers.pipeline(
+    "text-generation",
+    model=model,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+
+sequences = pipeline(
+    'I liked "Breaking Bad" and "Band of Brothers". Do you have any recommendations of other shows I might like?\n',
+    do_sample=True,
+    top_k=10,
+    num_return_sequences=1,
+    eos_token_id=tokenizer.eos_token_id,
+    max_length=200,
+)
+for seq in sequences:
+    print(f"Result: {seq['generated_text']}")