Update app.py

app.py

@@ -1,151 +1,120 @@
-import streamlit as st
-import openai
-import re
-import csv
-import base64
-from io import StringIO
-import threading
-from queue import Queue
-
-st.title("EinfachChatProjekt")
-
-api_key = st.sidebar.text_input("API Key:", value="sk-")
-openai.api_key = api_key
-
-show_notes = st.sidebar.checkbox("Show Notes", value="TRUE")
-data_section = st.sidebar.text_area("CSV or Text Data:")
-paste_data = st.sidebar.button("Paste Data")
-num_concurrent_calls = st.sidebar.number_input("Concurrent Calls:", min_value=1, max_value=2000, value=50, step=1)
-generate_all = st.sidebar.button("Generate All")
-reset = st.sidebar.button("Reset")
-add_row = st.sidebar.button("Add row")
-model = st.sidebar.selectbox("Model:", ["gpt-4", "gpt-3.5-turbo"])
-temperature = st.sidebar.slider("Temperature:", 0.0, 1.0, 0.6, step=0.01)
-max_tokens = st.sidebar.number_input("Max Tokens:", min_value=1, max_value=8192, value=2000, step=1)
-top_p = st.sidebar.slider("Top P:", 0.0, 1.0, 1.0, step=0.01)
-system_message = st.sidebar.text_area("System Message:")
-row_count = st.session_state.get("row_count", 1)
-
-if add_row:
-    row_count += 1
-    st.session_state.row_count = row_count
-
-if paste_data:
-    data = StringIO(data_section.strip())
-    reader = csv.reader(data, delimiter='\n', quotechar='"')
-    messages = [row[0] for row in reader]
-    if show_notes:
-        row_count = len(messages) // 2
-        for i in range(row_count):
-            st.session_state[f"note{i}"] = messages[i * 2]
-            st.session_state[f"message{i}"] = messages[i * 2 + 1]
-    else:
-        row_count = len(messages)
-        for i, message in enumerate(messages):
-            st.session_state[f"message{i}"] = message
-    st.session_state.row_count = row_count
-
-if reset:
-    row_count = 1
-    st.session_state.row_count = row_count
-    for i in range(100): # Assuming a maximum of 100 rows
-        st.session_state[f"note{i}"] = ""
-        st.session_state[f"message{i}"] = ""
-        st.session_state[f"response{i}"] = ""
-        st.session_state[f"prompt_tokens{i}"] = 0
-        st.session_state[f"response_tokens{i}"] = 0
-        st.session_state[f"word_count{i}"] = 0
-
-def generate_response(i, message):
-    try:
-        completion = openai.ChatCompletion.create(
-            model=model,
-            messages=[
-                {"role": "system", "content": system_message},
-                {"role": "user", "content": message}
-            ],
-            temperature=temperature,
-            max_tokens=max_tokens,
-            top_p=top_p
+from typing import Any, Dict, Tuple
+import warnings
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import (
+    StoppingCriteria,
+    StoppingCriteriaList,
+    TextIteratorStreamer,
+)
+
+
+INSTRUCTION_KEY = "### Instruction:"
+RESPONSE_KEY = "### Response:"
+END_KEY = "### End"
+INTRO_BLURB = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
+PROMPT_FOR_GENERATION_FORMAT = """{intro}
+{instruction_key}
+{instruction}
+{response_key}
+""".format(
+    intro=INTRO_BLURB,
+    instruction_key=INSTRUCTION_KEY,
+    instruction="{instruction}",
+    response_key=RESPONSE_KEY,
+)
+
+
+class InstructionTextGenerationPipeline:
+    def __init__(
+        self,
+        model_name,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        use_auth_token=None,
+    ) -> None:
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch_dtype,
+            trust_remote_code=trust_remote_code,
+            use_auth_token=use_auth_token,
         )
 
-        response = completion.choices[0].message.content
-        prompt_tokens = completion.usage['prompt_tokens']
-        response_tokens = completion.usage['total_tokens'] - prompt_tokens
-        word_count = len(re.findall(r'\w+', response))
-
-        return (i, response, prompt_tokens, response_tokens, word_count)
-
-    except Exception as e:
-        return (i, str(e), 0, 0, 0)
-
-def worker(q, results):
-    for item in iter(q.get, None):
-        results.put(generate_response(*item))
-
-class WorkerThread(threading.Thread):
-    def __init__(self, input_queue, output_queue):
-        threading.Thread.__init__(self)
-        self.input_queue = input_queue
-        self.output_queue = output_queue
-        self.daemon = True
-
-    def run(self):
-        while True:
-            i, message = self.input_queue.get()
-            try:
-                result = generate_response(i, message)
-                self.output_queue.put(result)
-            finally:
-                self.input_queue.task_done()
-
-if generate_all:
-    jobs = Queue()
-    results = Queue()
-
-    workers = [WorkerThread(jobs, results) for _ in range(num_concurrent_calls)]
-
-    for worker in workers:
-        worker.start()
-
-    for i in range(row_count):
-        message = st.session_state.get(f"message{i}", "")
-        jobs.put((i, message))
-
-    jobs.join()
-
-    while not results.empty():
-        i, response, prompt_tokens, response_tokens, word_count = results.get()
-        st.session_state[f"response{i}"] = response
-        st.session_state[f"prompt_tokens{i}"] = prompt_tokens
-        st.session_state[f"response_tokens{i}"] = response_tokens
-        st.session_state[f"word_count{i}"] = word_count
-
-def create_download_link(text, filename):
-    b64 = base64.b64encode(text.encode()).decode()
-    href = f'<a href="data:file/txt;base64,{b64}" download="{filename}">Download {filename}</a>'
-    return href
-
-for i in range(row_count):
-    if show_notes:
-        st.text_input(f"Note {i + 1}:", key=f"note{i}", value=st.session_state.get(f"note{i}", ""))
-    col1, col2 = st.columns(2)
-
-    with col1:
-        message = st.text_area(f"Message {i + 1}:", key=f"message{i}", value=st.session_state.get(f"message{i}", ""))
-
-        if st.button(f"Generate Response {i + 1}") and not st.session_state.get(f"response{i}", ""):
-            response, prompt_tokens, response_tokens, word_count = generate_response(i, message)
-            st.session_state[f"response{i}"] = response
-            st.session_state[f"prompt_tokens{i}"] = prompt_tokens
-            st.session_state[f"response_tokens{i}"] = response_tokens
-            st.session_state[f"word_count{i}"] = word_count
-
-    with col2:
-        st.text_area(f"Response {i + 1}:", value=st.session_state.get(f"response{i}", ""))
-        st.write(f"Tokens: {st.session_state.get(f'prompt_tokens{i}', 0)} / {st.session_state.get(f'response_tokens{i}', 0)} + Words: {st.session_state.get(f'word_count{i}', 0)}")
-
-responses_text = "\n\n".join([f"{st.session_state.get(f'note{i}', '')}\n{st.session_state.get(f'response{i}', '')}" for i in range(row_count) if show_notes] + [st.session_state.get(f"response{i}", "") for i in range(row_count) if not show_notes])
-download_filename = "GPT-4 Responses.txt"
-download_link = create_download_link(responses_text, download_filename)
-st.markdown(download_link, unsafe_allow_html=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=trust_remote_code,
+            use_auth_token=use_auth_token,
+        )
+        if tokenizer.pad_token_id is None:
+            warnings.warn(
+                "pad_token_id is not set for the tokenizer. Using eos_token_id as pad_token_id."
+            )
+            tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.padding_side = "left"
+        self.tokenizer = tokenizer
+
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.eval()
+        self.model.to(device=device, dtype=torch_dtype)
+
+        self.generate_kwargs = {
+            "temperature": 0.1,
+            "top_p": 0.92,
+            "top_k": 0,
+            "max_new_tokens": 1024,
+            "use_cache": True,
+            "do_sample": True,
+            "eos_token_id": self.tokenizer.eos_token_id,
+            "pad_token_id": self.tokenizer.pad_token_id,
+            "repetition_penalty": 1.1,  # 1.0 means no penalty, > 1.0 means penalty, 1.2 from CTRL paper
+        }
+
+    def format_instruction(self, instruction):
+        return PROMPT_FOR_GENERATION_FORMAT.format(instruction=instruction)
+
+    def __call__(
+        self, instruction: str, **generate_kwargs: Dict[str, Any]
+    ) -> Tuple[str, str, float]:
+        s = PROMPT_FOR_GENERATION_FORMAT.format(instruction=instruction)
+        input_ids = self.tokenizer(s, return_tensors="pt").input_ids
+        input_ids = input_ids.to(self.model.device)
+        gkw = {**self.generate_kwargs, **generate_kwargs}
+        with torch.no_grad():
+            output_ids = self.model.generate(input_ids, **gkw)
+        # Slice the output_ids tensor to get only new tokens
+        new_tokens = output_ids[0, len(input_ids[0]) :]
+        output_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
+        return output_text
+
+# Initialize the model and tokenizer
+generate = InstructionTextGenerationPipeline(
+    "mosaicml/mpt-7b-instruct",
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+)
+stop_token_ids = generate.tokenizer.convert_tokens_to_ids(["<|endoftext|>"])
+
+
+# Define a custom stopping criteria
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        for stop_id in stop_token_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+"""### The prompt & response"""
+
+import json
+import textwrap
+
+def get_prompt(instruction):
+    prompt_template = f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:"
+    return prompt_template
+
+# print(get_prompt('What is the meaning of life?'))
+
+def parse_text(text):
+        wrapped_text = textwrap.fill(text, width=100)
+        print(wrapped_text +'\n\n')