Update app.py

app.py

@@ -1,21 +1,33 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
-device = "cuda" # the device to load the model onto
+from transformers import AutoTokenizer, AutoModelWithLMHead  
+import torch
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = "cpu"
+    
+    
+tokenizer = AutoTokenizer.from_pretrained("salesken/content_generation_from_phrases") 
+model = AutoModelWithLMHead.from_pretrained("salesken/content_generation_from_phrases").to(device)
 
-model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
 
-messages = [
-    {"role": "user", "content": "What is your favourite condiment?"},
-    {"role": "assistant", "content": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!"},
-    {"role": "user", "content": "Do you have mayonnaise recipes?"}
-]
+input_query=["data science beginner"]
+query = "<|startoftext|> " + input_query[0] + " ~~"
 
-encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
+input_ids = tokenizer.encode(query.lower(), return_tensors='pt').to(device)
+sample_outputs = model.generate(input_ids,
+                                do_sample=True,
+                                num_beams=1, 
+                                max_length=256,
+                                temperature=0.9,
+                                top_k = 30,
+                                num_return_sequences=100)
+content = []
+for i in range(len(sample_outputs)):
+    r = tokenizer.decode(sample_outputs[i], skip_special_tokens=True).split('||')[0]
+    r = r.split(' ~~ ')[1]
+    if r not in content:
+        content.append(r)
 
-model_inputs = encodeds.to(device)
-model.to(device)
-
-generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-decoded = tokenizer.batch_decode(generated_ids)
-st.write(decoded[0])
+st.write(content)