Update app.py

app.py

@@ -1,11 +1,7 @@
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 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"
@@ -16,7 +12,7 @@ eng_trans_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled
 translator = pipeline('translation', model=trans_model, tokenizer=eng_trans_tokenizer, src_lang="eng_Latn", tgt_lang='sin_Sinh', max_length=400, device=device)
 
 sin_trans_model = AutoModelForSeq2SeqLM.from_pretrained("thilina/mt5-sinhalese-english").to(device)
-si_trans_tokenizer = AutoTokenizer.from_pretrained("thilina/mt5-sinhalese-english", use_fast=False)  # Use slow tokenizer
+si_trans_tokenizer = AutoTokenizer.from_pretrained("thilina/mt5-sinhalese-english", use_fast=False)
 
 singlish_pipe = pipeline("text2text-generation", model="Dhahlan2000/Simple_Translation-model-for-GPT-v14")
 
@@ -47,64 +43,11 @@ def transliterate_from_sinhala(text):
 def transliterate_to_sinhala(text):
     return transliterate.process('Velthuis', 'Sinhala', text)
 
-# Load conversation model
-# 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"
-
-# tokenizer = AutoTokenizer.from_pretrained(model)
-# text_gen_pipeline = pipeline(
-#     "text-generation",
-#     model=model,
-#     tokenizer=tokenizer,
-#     torch_dtype=torch.bfloat16,
-#     trust_remote_code=True,
-#     device_map="auto",
-# )
-
-# pipe1 = pipeline("text-generation", model="unsloth/gemma-2b-it")
-
-# client = InferenceClient("google/gemma-2b-it")
+# Placeholder for conversation model loading and pipeline setup
 
 def conversation_predict(text):
-    interface = gr.interface.load("microsoft/Phi-3-mini-4k-instruct")
+    interface = gr.Interface.load("microsoft/Phi-3-mini-4k-instruct")
     return interface([text])[0]
-    # return client.text_generation(text, return_full_text=False)
-    # pipe = pipeline(
-    # "text-generation",
-    # model=model,
-    # tokenizer=tokenizer,
-    # )
-    # generation_args = {
-    # "max_new_tokens": 500,
-    # "return_full_text": False,
-    # "temperature": 0.0,
-    # "do_sample": False,
-    # }
-
-    # output = pipe(text, **generation_args)
-    # return output[0]['generated_text']
-    # input_ids = tokenizer(text, return_tensors="pt")
-    # 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"]
-
-    # sequences = text_gen_pipeline(
-    #     text,
-    #     max_length=200,
-    #     do_sample=True,
-    #     top_k=10,
-    #     num_return_sequences=1,
-    #     eos_token_id=tokenizer.eos_token_id,
-    # )
-    # return sequences[0]['generated_text']
-    
-
 
 def ai_predicted(user_input):
     if user_input.lower() == 'exit':
@@ -141,7 +84,6 @@ def respond(
     messages.append({"role": "user", "content": message})
 
     response = ai_predicted(message)
-    # response = pipe1({"role": "user", "content": message})
 
     yield response