Update the logic to comply with Llama 2

app.py

@@ -1,28 +1,33 @@
 import os
-import requests
-import gradio as gr
+import torch
 from dotenv import load_dotenv, find_dotenv
-from transformers import pipeline
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 
 # Load environment variables
 _ = load_dotenv(find_dotenv()) 
 hf_api_key = os.environ['HF_API_KEY']
 
 class Client:
-    def __init__(self, base_url, headers, timeout):
-        self.base_url = base_url
-        self.headers = headers
-        self.timeout = timeout
-
-    def generate(self, prompt, max_new_tokens):
-        # Your text generation code here.
-        return {'generated_text': 'Placeholder text'}
-      
+    def __init__(self, model_name, device_map="auto"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map=device_map,
+            torch_dtype=torch.float16,
+            load_in_8bit=True,
+            rope_scaling={"type": "dynamic", "factor": 2}
+        )
+
     def generate_stream(self, prompt, max_new_tokens, stop_sequences, temperature):
-        # Your streaming text generation code here.
-        pass
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        del inputs["token_type_ids"]
+        streamer = TextStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+        output = self.model.generate(**inputs, streamer=streamer, use_cache=True, max_new_tokens=float('inf'))
+        output_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return output_text
 
-client = Client(os.environ['HF_API_FALCOM_BASE'], headers={"Authorization": f"Basic {hf_api_key}"}, timeout=120)
+client = Client("upstage/Llama-2-70b-instruct")
 
 def format_chat_prompt(message, chat_history, instruction):
     prompt = f"System:{instruction}"
@@ -35,20 +40,11 @@ def format_chat_prompt(message, chat_history, instruction):
 def respond(message, chat_history, instruction, temperature=0.7):
     prompt = format_chat_prompt(message, chat_history, instruction)
     chat_history = chat_history + [[message, ""]]
-    stream = client.generate_stream(prompt, max_new_tokens=1024, stop_sequences=["\nUser:", ""], temperature=temperature)
-    acc_text = ""
-    for idx, response in enumerate(stream):
-            text_token = response.token.text
-            if response.details:
-                return
-            if idx == 0 and text_token.startswith(" "):
-                text_token = text_token[1:]
-            acc_text += text_token
-            last_turn = list(chat_history.pop(-1))
-            last_turn[-1] += acc_text
-            chat_history = chat_history + [last_turn]
-            yield "", chat_history
-            acc_text = ""
+    output_text = client.generate_stream(prompt, max_new_tokens=1024, stop_sequences=["\nUser:", ""], temperature=temperature)
+    last_turn = list(chat_history.pop(-1))
+    last_turn[-1] += output_text
+    chat_history = chat_history + [last_turn]
+    return "", chat_history
 
 iface = gr.Interface(fn=respond, inputs=[gr.Textbox(label="Prompt"), gr.Chatbot(label="Chat History", height=240), gr.Textbox(label="System message", lines=2, value="A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers."), gr.Slider(label="temperature", minimum=0.1, maximum=1, value=0.7, step=0.1)], outputs=[gr.Textbox(label="Prompt"), gr.Chatbot(label="Chat History", height=240)])