Update app.py

app.py

@@ -3,6 +3,9 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoMod
 from aksharamukha import transliterate
 import torch
 from huggingface_hub import InferenceClient
+import os
+
+access_token = os.environ["TOKEN"]
 
 # Set up device
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -45,9 +48,9 @@ def transliterate_to_sinhala(text):
     return transliterate.process('Velthuis', 'Sinhala', text)
 
 # Load conversation model
-# conv_model_name = "microsoft/Phi-3-mini-4k-instruct"  # Use GPT-2 instead of the gated model
-# tokenizer = AutoTokenizer.from_pretrained(conv_model_name, trust_remote_code=True)
-# model = AutoModelForCausalLM.from_pretrained(conv_model_name, trust_remote_code=True).to(device)
+conv_model_name = "google/gemma-2b-it"  # Use GPT-2 instead of the gated model
+tokenizer = AutoTokenizer.from_pretrained(conv_model_name, trust_remote_code=True, token = access_token)
+model = AutoModelForCausalLM.from_pretrained(conv_model_name, trust_remote_code=True, token = access_token, torch_dtype=torch.bfloat16).to(device)
 # pipe1 = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0").to(device)
 
 # model = "tiiuae/falcon-7b-instruct"
@@ -64,10 +67,10 @@ def transliterate_to_sinhala(text):
 
 # pipe1 = pipeline("text-generation", model="unsloth/gemma-2b-it")
 
-client = InferenceClient("google/gemma-2b-it")
+# client = InferenceClient("google/gemma-2b-it")
 
 # def conversation_predict(text):
-    return client.text_generation(text, return_full_text=False)
+    # return client.text_generation(text, return_full_text=False)
     # pipe = pipeline(
     # "text-generation",
     # model=model,
@@ -82,9 +85,9 @@ client = InferenceClient("google/gemma-2b-it")
 
     # output = pipe(text, **generation_args)
     # return output[0]['generated_text']
-    # input_ids = tokenizer(text, return_tensors="pt").to(device)
-    # outputs = model.generate(**input_ids)
-    # return tokenizer.decode(outputs[0])
+    input_ids = tokenizer(text, return_tensors="pt").to(device)
+    outputs = model.generate(**input_ids)
+    return tokenizer.decode(outputs[0])
 
     # outputs = pipe1(text, max_new_tokens=256, temperature=0.7, top_k=50, top_p=0.95)
     # return outputs[0]["generated_text"]