Update app.py

Upgrade model to 7B version

app.py

@@ -30,7 +30,34 @@ wiki_wiki = wikipediaapi.Wikipedia('Organika ([email protected])', 'en')
 # Use a pipeline as a high-level helper
 from transformers import pipeline
 topic_model = pipeline("zero-shot-classification", model="valhalla/distilbart-mnli-12-9", device=0)
-model = pipeline("text-generation", model="Organika/StarCoder-1B-WoW-JSON", device=0)
+#model = pipeline("text-generation", model="Organika/StarCoder-7B-WoW-JSON_1", device=0)
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+bnb_config = BitsAndBytesConfig(
+load_in_4bit=True,
+bnb_4bit_use_double_quant=True,
+bnb_4bit_quant_type="nf4",
+bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+model_name = "umm-maybe/StarCoder-7B-WoW-JSON_1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map="auto")
+
+def generate_text(prompt):
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(
+        inputs,
+        do_sample=True,
+        max_new_tokens=200,
+        temperature=0.6,
+        top_p=0.9,
+        top_k=40,
+        repetition_penalty=1.1
+        )
+    results = self.tokenizer.decode(outputs[0], clean_up_tokenization_spaces=False) #.replace(prompt,"")
+    return results
 
 def merlin_chat(message, history):
     chat_text = ""
@@ -89,7 +116,8 @@ def merlin_chat(message, history):
     print(f"PROMPT: {prompt}")
     for attempt in range(3):
         #result = generate_text(prompt, model_path, parameters, headers)
-        result = model(prompt,return_full_text=False, max_new_tokens=256, temperature=0.8, repetition_penalty=1.1)
+        #result = model(prompt,return_full_text=False, max_new_tokens=256, temperature=0.8, repetition_penalty=1.1)
+        result = generate_text(prompt)
         response = result[0]['generated_text']
         print(f"COMPLETION: {response}") # so we can see it in logs
         start = 0
@@ -102,7 +130,7 @@ def merlin_chat(message, history):
         if end<=0:
             continue
         cleanStr = cleanStr[:end]
-        messageStr = prompt + cleanStr + ']'
+        messageStr = cleanStr + ']'
         messages = json.loads(messageStr)
         message = messages[-1]
         if message['role'] != 'assistant':