Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
-
 import torch 
 import torch.nn as nn
 import torch.nn.functional as FU
 from torchvision import transforms
 import numpy as np 
 from ultralytics import YOLO
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
+import os
 import torch
 from pypinyin import pinyin
 from PIL import Image, ImageDraw, ImageEnhance
@@ -263,20 +264,21 @@ test_transforms = transforms.Compose([
 
 
 """import llm from huggingface"""
-model_name = "IncreasingLoss/FineTunedTranslation_Yi-1.5-9B-Chat-Q6_K/"
-# Yi-1.5-9B-Chat-Q6_K.gguf
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-hf_model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto" if torch.cuda.is_available() else None,
-    trust_remote_code=True
+model_path = hf_hub_download(
+    repo_id="IncreasingLoss/FineTunedTranslation_Yi-1.5-9B-Chat-Q6_K",
+    filename="Yi-1.5-9B-Chat-Q6_K.gguf",
+    local_dir="./models",  # Optional: specify where to save
+    local_dir_use_symlinks=False
 )
 
-# Set padding token if not already set
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+print(f"Model downloaded to: {model_path}")
+
+llm = Llama(
+    model_path=model_path,
+    n_ctx=4096,
+    n_gpu_layers=-1,  # Use all GPU layers
+    verbose=False
+)
 
 """app functions"""
 def convert_to_pil(inputimg):
@@ -436,10 +438,11 @@ def generate_translation(prompt, max_new_tokens=512):
     response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
     return response.strip()
 
-# Updated translate_text function
+
 def translate_text(threshhold_slider, is_vertical, input_img):
     # convert input to pil image
     pil_img = convert_to_pil(input_img)
+    # FIX #2: Add null check for pil_img
     if pil_img is None:
         return "", "", ""
 
@@ -449,16 +452,16 @@ def translate_text(threshhold_slider, is_vertical, input_img):
         chinese_text = ""
         for crop in sorted_cropped_images:
             crop = test_transforms(crop)
-            crop = crop.to(user_device)
-            crop_batch = crop.unsqueeze(dim=0)
-            
+            crop = crop.to(user_device)  # Move tensor to device
+            crop_batch = crop.unsqueeze(dim=0)  # Add batch dimension
+
             classifier.to(user_device)
             
             logits = classifier(crop_batch)
             class_idx = logits.argmax(dim=1)
             
-            class_idx_cpu = class_idx.cpu().item()
-            chinese_charcter = classes_dict[class_idx_cpu]
+            class_idx_cpu = class_idx.cpu().item()  # Convert to CPU integer
+            chinese_charcter = classes_dict[class_idx_cpu]  # Use integer index for list
             chinese_text += chinese_charcter
 
     pinyin_text = pinyin(chinese_text)
@@ -466,24 +469,34 @@ def translate_text(threshhold_slider, is_vertical, input_img):
     for pin in pinyin_text:
         pinyin_sentence += f"{pin[0]} "
 
-    # Enhanced prompt for better translation quality
-    prompt = f"""You are a professional Chinese to English translator. 
-1. Translate the following Chinese text to natural, fluent English:
-"{chinese_text}"
-2. Respond only with the translated English text.
-English translation: """
 
-    # Use Hugging Face model instead of llama_cpp
-    translation = generate_translation(prompt, max_new_tokens=min(len(chinese_text)*2, 512))
+    prompt = f""" You are a professional Chinese to English translator. 
+        1. Translate the following Chinese text to natural, fluent English:
+        "{chinese_text}"
+        2. Respond only with the translated English text.
+        English translation: """
+
+    output = llm(
+        prompt, 
+        max_tokens=min(len(chinese_text)*2, 512),  # More reasonable token limit with safety cap
+        stop=["</s>", "\n\n", "Chinese:", "Chinese text:"],  # Enhanced stop tokens
+        echo=False,
+        temperature=0.3,  # Slightly higher temperature for more natural language
+        frequency_penalty=.5,  # Penalize repetition of tokens - higher values discourage repetition
+        presence_penalty=.5,  # Penalize tokens that have appeared at all - prevents topic repetition
+        top_p=0.9,  # Nucleus sampling - restricts to top 90% most likely tokens for more coherent output
+        stream=False  # Explicitly disable streaming to get a dict response
+    )
     
-    # Clean up the translation
+    # Extract the translation
+    translation = output["choices"][0]["text"].strip()
     try:
-        if '"' in translation:
-            start = translation.index('"') + 1
-            end = translation.rindex('"')
-            translation = translation[start:end]
+        start = translation.index('"') + 1
+        # find index of last double-quote
+        end = translation.rindex('"')
+        translation = translation[start:end]    
     except:
-        pass
+        pass    
 
     return chinese_text, pinyin_sentence, translation