Update app.py

app.py

@@ -1,28 +1,34 @@
 import streamlit as st
 import time
-from transformers import AutoModelForCausalLM, AutoTokenizer
+#from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-@st.cache(allow_output_mutation=True)
-def define_model():
-    model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype=torch.float16).cuda()
-    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b", use_fast=False)
-    return model, tokenizer
+#@st.cache(allow_output_mutation=True)
+#def define_model():
+#    model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype=torch.float16).cuda()
+#    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b", use_fast=False)
+#    return model, tokenizer
+    
+from transformers import pipeline
 
-@st.cache(allow_output_mutation=True)
-def opt_model(prompt, model, tokenizer, num_sequences = 1, max_length = 50):  
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
-    generated_ids = model.generate(input_ids, num_return_sequences=num_sequences, max_length=max_length)
-    answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-    return answer
+generator = pipeline('text-generation', model="facebook/opt-1.3b")
+answer = generator("Hello, I'm am conscious and")
 
 
-model, tokenizer = define_model()
+#@st.cache(allow_output_mutation=True)
+#def opt_model(prompt, model, tokenizer, num_sequences = 1, max_length = 50):  
+#    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
+#    generated_ids = model.generate(input_ids, num_return_sequences=num_sequences, max_length=max_length)
+#    answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+#    return answer
+
+
+#model, tokenizer = define_model()
 
 prompt= st.text_area('Your prompt here',
  '''Hello, I'm am conscious and''')
  
-answer = opt_model(prompt, model, tokenizer,)
+#answer = opt_model(prompt, model, tokenizer,)
 #lst = ['ciao come stai sjfsbd dfhsdf  fuahfuf  feuhfu wefwu ']
 lst = ' '.join(answer)